Alerting: Change screenshots to use components (#55156)

* Alerting: Change screenshots to use components

This commit changes screenshots to use a number of components instead of a set of functional wrappers.

It moves the uploading of screenshots from the screenshot package to the image package so we can re-use the same code for both uploading screenshots and server-side images; SingleFlight from the screenshot package to the image package so we can use it for both taking and uploading the screenshot, where as before it was used just for taking the screenshot; and it also removes the use of a cache because we know that screenshots can be taken at most once per tick of the scheduler.

pkg/services/ngalert/image/service.go

@@ -7,10 +7,12 @@ import (
 	"time"
 
 	"github.com/prometheus/client_golang/prometheus"
+	"golang.org/x/sync/singleflight"
 
 	"github.com/grafana/grafana/pkg/components/imguploader"
+	"github.com/grafana/grafana/pkg/infra/log"
 	"github.com/grafana/grafana/pkg/services/dashboards"
-	ngmodels "github.com/grafana/grafana/pkg/services/ngalert/models"
+	"github.com/grafana/grafana/pkg/services/ngalert/models"
 	"github.com/grafana/grafana/pkg/services/ngalert/store"
 	"github.com/grafana/grafana/pkg/services/rendering"
 	"github.com/grafana/grafana/pkg/services/screenshot"
@@ -48,7 +50,7 @@ func ProvideDeleteExpiredService(store *store.DBstore) *DeleteExpiredService {
 //go:generate mockgen -destination=mock.go -package=image github.com/grafana/grafana/pkg/services/ngalert/image ImageService
 type ImageService interface {
 	// NewImage returns a new image for the alert instance.
-	NewImage(ctx context.Context, r *ngmodels.AlertRule) (*ngmodels.Image, error)
+	NewImage(ctx context.Context, r *models.AlertRule) (*models.Image, error)
 }
 
 // ScreenshotImageService takes screenshots of the alert rule and saves the
@@ -56,45 +58,56 @@ type ImageService interface {
 // as an annotation or label to the Alertmanager. This service cannot take
 // screenshots of alert rules that are not associated with a dashboard panel.
 type ScreenshotImageService struct {
+	limiter      screenshot.RateLimiter
+	logger       log.Logger
 	screenshots  screenshot.ScreenshotService
+	singleflight singleflight.Group
 	store        store.ImageStore
+	uploads      *UploadingService
 }
 
-func NewScreenshotImageService(screenshots screenshot.ScreenshotService, store store.ImageStore) ImageService {
+// NewScreenshotImageService returns a new ScreenshotImageService.
+func NewScreenshotImageService(
+	limiter screenshot.RateLimiter,
+	logger log.Logger,
+	screenshots screenshot.ScreenshotService,
+	store store.ImageStore,
+	uploads *UploadingService) ImageService {
 	return &ScreenshotImageService{
+		limiter:     limiter,
+		logger:      logger,
 		screenshots: screenshots,
 		store:       store,
+		uploads:     uploads,
 	}
 }
 
 // NewScreenshotImageServiceFromCfg returns a new ScreenshotImageService
 // from the configuration.
-func NewScreenshotImageServiceFromCfg(cfg *setting.Cfg, metrics prometheus.Registerer,
-	db *store.DBstore, ds dashboards.DashboardService, rs rendering.Service) (ImageService, error) {
-	// If screenshots are disabled then return the ScreenshotUnavailableService
-	if !cfg.UnifiedAlerting.Screenshots.Capture {
-		return &ScreenshotImageService{
-			screenshots: &screenshot.ScreenshotUnavailableService{},
-		}, nil
-	}
+func NewScreenshotImageServiceFromCfg(cfg *setting.Cfg, db *store.DBstore, ds dashboards.DashboardService,
+	rs rendering.Service, r prometheus.Registerer) (ImageService, error) {
+	var (
+		limiter     screenshot.RateLimiter       = &screenshot.NoOpRateLimiter{}
+		screenshots screenshot.ScreenshotService = &screenshot.ScreenshotUnavailableService{}
+		uploads     *UploadingService            = nil
+	)
 
-	// Image uploading is an optional feature of screenshots
-	s := screenshot.NewRemoteRenderScreenshotService(ds, rs)
+	// If screenshots are enabled
+	if cfg.UnifiedAlerting.Screenshots.Capture {
+		limiter = screenshot.NewTokenRateLimiter(cfg.UnifiedAlerting.Screenshots.MaxConcurrentScreenshots)
+		screenshots = screenshot.NewHeadlessScreenshotService(ds, rs, r)
+
+		// Image uploading is an optional feature
 		if cfg.UnifiedAlerting.Screenshots.UploadExternalImageStorage {
-		u, err := imguploader.NewImageUploader()
+			m, err := imguploader.NewImageUploader()
 			if err != nil {
 				return nil, fmt.Errorf("failed to initialize uploading screenshot service: %w", err)
 			}
-		s = screenshot.NewUploadingScreenshotService(metrics, s, u)
+			uploads = NewUploadingService(m, r)
 		}
-	s = screenshot.NewRateLimitScreenshotService(s, cfg.UnifiedAlerting.Screenshots.MaxConcurrentScreenshots)
-	s = screenshot.NewSingleFlightScreenshotService(s)
-	s = screenshot.NewCachableScreenshotService(metrics, screenshotCacheTTL, s)
-	s = screenshot.NewObservableScreenshotService(metrics, s)
-	return &ScreenshotImageService{
-		store:       db,
-		screenshots: s,
-	}, nil
+	}
+
+	return NewScreenshotImageService(limiter, cfg.Logger, screenshots, db, uploads), nil
 }
 
 // NewImage returns a screenshot of the alert rule or an error.
@@ -104,10 +117,11 @@ func NewScreenshotImageServiceFromCfg(cfg *setting.Cfg, metrics prometheus.Regis
 // or the dashboard does not exist, an ErrNoDashboard error is returned. If the
 // alert rule has a Dashboard UID and the dashboard exists, but does not have a
 // Panel ID in its annotations then an ErrNoPanel error is returned.
-func (s *ScreenshotImageService) NewImage(ctx context.Context, r *ngmodels.AlertRule) (*ngmodels.Image, error) {
+func (s *ScreenshotImageService) NewImage(ctx context.Context, r *models.AlertRule) (*models.Image, error) {
 	if r.DashboardUID == nil {
 		return nil, ErrNoDashboard
 	}
+
 	if r.PanelID == nil || *r.PanelID == 0 {
 		return nil, ErrNoPanel
 	}
@@ -115,41 +129,51 @@ func (s *ScreenshotImageService) NewImage(ctx context.Context, r *ngmodels.Alert
 	ctx, cancelFunc := context.WithTimeout(ctx, screenshotTimeout)
 	defer cancelFunc()
 
-	screenshot, err := s.screenshots.Take(ctx, screenshot.ScreenshotOptions{
-		Timeout:      screenshotTimeout,
+	opts := screenshot.ScreenshotOptions{
 		DashboardUID: *r.DashboardUID,
 		PanelID:      *r.PanelID,
-	})
+		Timeout:      screenshotTimeout,
+	}
+
+	//
+	k := fmt.Sprintf("%s-%d-%s", opts.DashboardUID, opts.PanelID, opts.Theme)
+	result, err, _ := s.singleflight.Do(k, func() (interface{}, error) {
+		screenshot, err := s.limiter.Do(ctx, opts, s.screenshots.Take)
 		if err != nil {
-		// TODO: Check for screenshot upload failures. These images should still be
-		// stored because we have a local disk path that could be useful.
 			if errors.Is(err, dashboards.ErrDashboardNotFound) {
 				return nil, ErrNoDashboard
 			}
 			return nil, err
 		}
-
-	v := ngmodels.Image{
-		Path: screenshot.Path,
-		URL:  screenshot.URL,
+		image := models.Image{Path: screenshot.Path}
+		if s.uploads != nil {
+			if image, err = s.uploads.Upload(ctx, image); err != nil {
+				s.logger.Warn("failed to upload image", "path", image.Path, "err", err)
 			}
-	if err := s.store.SaveImage(ctx, &v); err != nil {
+		}
+		if err := s.store.SaveImage(ctx, &image); err != nil {
 			return nil, fmt.Errorf("failed to save image: %w", err)
 		}
+		return image, nil
+	})
+	if err != nil {
+		return nil, err
+	}
 
-	return &v, nil
+	image := result.(models.Image)
+	return &image, nil
 }
 
 // NotAvailableImageService is a service that returns ErrScreenshotsUnavailable.
 type NotAvailableImageService struct{}
 
-func (s *NotAvailableImageService) NewImage(_ context.Context, _ *ngmodels.AlertRule) (*ngmodels.Image, error) {
+func (s *NotAvailableImageService) NewImage(_ context.Context, _ *models.AlertRule) (*models.Image, error) {
 	return nil, screenshot.ErrScreenshotsUnavailable
 }
 
 // NoopImageService is a no-op image service.
 type NoopImageService struct{}
 
-func (s *NoopImageService) NewImage(_ context.Context, _ *ngmodels.AlertRule) (*ngmodels.Image, error) {
-	return &ngmodels.Image{}, nil
+func (s *NoopImageService) NewImage(_ context.Context, _ *models.AlertRule) (*models.Image, error) {
+	return &models.Image{}, nil
 }

pkg/services/ngalert/image/service_test.go

@@ -0,0 +1,89 @@
+package image
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/golang/mock/gomock"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"github.com/xorcare/pointer"
+
+	"github.com/grafana/grafana/pkg/components/imguploader"
+	"github.com/grafana/grafana/pkg/infra/log"
+	"github.com/grafana/grafana/pkg/services/ngalert/models"
+	"github.com/grafana/grafana/pkg/services/ngalert/store"
+	"github.com/grafana/grafana/pkg/services/screenshot"
+)
+
+func TestScreenshotImageService(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	var (
+		images      = store.NewFakeImageStore(t)
+		limiter     = screenshot.NoOpRateLimiter{}
+		screenshots = screenshot.NewMockScreenshotService(ctrl)
+		uploads     = imguploader.NewMockImageUploader(ctrl)
+	)
+
+	s := NewScreenshotImageService(&limiter, log.NewNopLogger(), screenshots, images,
+		NewUploadingService(uploads, prometheus.NewRegistry()))
+
+	ctx := context.Background()
+
+	// assert that a screenshot is taken
+	screenshots.EXPECT().Take(gomock.Any(), screenshot.ScreenshotOptions{
+		DashboardUID: "foo",
+		PanelID:      1,
+		Timeout:      screenshotTimeout,
+	}).Return(&screenshot.Screenshot{
+		Path: "foo.png",
+	}, nil)
+
+	// the screenshot is made into an image and uploaded
+	uploads.EXPECT().Upload(gomock.Any(), "foo.png").
+		Return("https://example.com/foo.png", nil)
+
+	// and then saved into the database
+	expected := models.Image{
+		ID:    1,
+		Token: "foo",
+		Path:  "foo.png",
+		URL:   "https://example.com/foo.png",
+	}
+
+	image, err := s.NewImage(ctx, &models.AlertRule{
+		DashboardUID: pointer.String("foo"),
+		PanelID:      pointer.Int64(1)})
+	require.NoError(t, err)
+	assert.Equal(t, expected, *image)
+
+	// assert that a screenshot is taken
+	screenshots.EXPECT().Take(gomock.Any(), screenshot.ScreenshotOptions{
+		DashboardUID: "bar",
+		PanelID:      1,
+		Timeout:      screenshotTimeout,
+	}).Return(&screenshot.Screenshot{
+		Path: "bar.png",
+	}, nil)
+
+	// the screenshot is made into an image and uploaded, but the upload returns an error
+	uploads.EXPECT().Upload(gomock.Any(), "bar.png").
+		Return("", errors.New("failed to upload bar.png"))
+
+	// and then saved into the database, but without a URL
+	expected = models.Image{
+		ID:    2,
+		Token: "bar",
+		Path:  "bar.png",
+	}
+
+	image, err = s.NewImage(ctx, &models.AlertRule{
+		DashboardUID: pointer.String("bar"),
+		PanelID:      pointer.Int64(1)})
+	require.NoError(t, err)
+	assert.Equal(t, expected, *image)
+}

pkg/services/ngalert/image/upload.go

@@ -0,0 +1,47 @@
+package image
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promauto"
+
+	"github.com/grafana/grafana/pkg/components/imguploader"
+	ngmodels "github.com/grafana/grafana/pkg/services/ngalert/models"
+)
+
+type UploadingService struct {
+	uploader  imguploader.ImageUploader
+	failures  prometheus.Counter
+	successes prometheus.Counter
+}
+
+func NewUploadingService(uploader imguploader.ImageUploader, r prometheus.Registerer) *UploadingService {
+	return &UploadingService{
+		uploader: uploader,
+		failures: promauto.With(r).NewCounter(prometheus.CounterOpts{
+			Name:      "image_upload_failures_total",
+			Namespace: "grafana",
+			Subsystem: "alerting",
+		}),
+		successes: promauto.With(r).NewCounter(prometheus.CounterOpts{
+			Name:      "image_upload_successes_total",
+			Namespace: "grafana",
+			Subsystem: "alerting",
+		}),
+	}
+}
+
+// Upload uploads an image and returns a new image with the unmodified path and a URL.
+// It returns the unmodified image on error.
+func (s *UploadingService) Upload(ctx context.Context, image ngmodels.Image) (ngmodels.Image, error) {
+	url, err := s.uploader.Upload(ctx, image.Path)
+	if err != nil {
+		defer s.failures.Inc()
+		return image, fmt.Errorf("failed to upload screenshot: %w", err)
+	}
+	image.URL = url
+	defer s.successes.Inc()
+	return image, nil
+}

pkg/services/ngalert/image/upload_test.go

@@ -0,0 +1,41 @@
+package image
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/golang/mock/gomock"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/grafana/grafana/pkg/components/imguploader"
+	"github.com/grafana/grafana/pkg/services/ngalert/models"
+)
+
+func TestUploadingService(t *testing.T) {
+	c := gomock.NewController(t)
+	defer c.Finish()
+
+	u := imguploader.NewMockImageUploader(c)
+	s := NewUploadingService(u, prometheus.NewRegistry())
+
+	ctx := context.Background()
+
+	u.EXPECT().Upload(ctx, "foo.png").Return("https://example.com/foo.png", nil)
+	image, err := s.Upload(ctx, models.Image{Path: "foo.png"})
+	require.NoError(t, err)
+	assert.Equal(t, models.Image{
+		Path: "foo.png",
+		URL:  "https://example.com/foo.png",
+	}, image)
+
+	// error on upload should still return screenshot on disk
+	u.EXPECT().Upload(ctx, "foo.png").Return("", errors.New("service is unavailable"))
+	image, err = s.Upload(ctx, models.Image{Path: "foo.png"})
+	assert.EqualError(t, err, "failed to upload screenshot: service is unavailable")
+	assert.Equal(t, models.Image{
+		Path: "foo.png",
+	}, image)
+}

pkg/services/ngalert/ngalert.go

@@ -128,7 +128,7 @@ func (ng *AlertNG) init() error {
 		return err
 	}
 
-	imageService, err := image.NewScreenshotImageServiceFromCfg(ng.Cfg, ng.Metrics.Registerer, store, ng.dashboardService, ng.renderService)
+	imageService, err := image.NewScreenshotImageServiceFromCfg(ng.Cfg, store, ng.dashboardService, ng.renderService, ng.Metrics.Registerer)
 	if err != nil {
 		return err
 	}

pkg/services/ngalert/store/testing.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"math/rand"
+	"strings"
 	"sync"
 	"testing"
 
@@ -14,6 +15,57 @@ import (
 	"github.com/grafana/grafana/pkg/services/ngalert/models"
 )
 
+func NewFakeImageStore(t *testing.T) *FakeImageStore {
+	return &FakeImageStore{
+		t:      t,
+		images: make(map[string]*models.Image),
+	}
+}
+
+type FakeImageStore struct {
+	t      *testing.T
+	mtx    sync.Mutex
+	images map[string]*models.Image
+}
+
+func (s *FakeImageStore) GetImage(_ context.Context, token string) (*models.Image, error) {
+	s.mtx.Lock()
+	defer s.mtx.Unlock()
+	if image, ok := s.images[token]; ok {
+		return image, nil
+	}
+	return nil, models.ErrImageNotFound
+}
+
+func (s *FakeImageStore) GetImages(_ context.Context, tokens []string) ([]models.Image, []string, error) {
+	s.mtx.Lock()
+	defer s.mtx.Unlock()
+	images := make([]models.Image, 0, len(tokens))
+	for _, token := range tokens {
+		if image, ok := s.images[token]; ok {
+			images = append(images, *image)
+		}
+	}
+	if len(images) < len(tokens) {
+		return images, unmatchedTokens(tokens, images), models.ErrImageNotFound
+	}
+	return images, nil, nil
+}
+
+func (s *FakeImageStore) SaveImage(_ context.Context, image *models.Image) error {
+	s.mtx.Lock()
+	defer s.mtx.Unlock()
+	if image.ID == 0 {
+		image.ID = int64(len(s.images)) + 1
+	}
+	if image.Token == "" {
+		tmp := strings.Split(image.Path, ".")
+		image.Token = strings.Join(tmp[:len(tmp)-1], ".")
+	}
+	s.images[image.Token] = image
+	return nil
+}
+
 func NewFakeRuleStore(t *testing.T) *FakeRuleStore {
 	return &FakeRuleStore{
 		t:     t,

pkg/services/screenshot/cache.go

@@ -0,0 +1,72 @@
+package screenshot
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	gocache "github.com/patrickmn/go-cache"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promauto"
+)
+
+// CacheService caches screenshots.
+//
+//go:generate mockgen -destination=cache_mock.go -package=screenshot github.com/grafana/grafana/pkg/services/screenshot CacheService
+type CacheService interface {
+	// Get returns the screenshot for the options or false if a screenshot with these
+	// options does not exist.
+	Get(ctx context.Context, opts ScreenshotOptions) (*Screenshot, bool)
+	// Set the screenshot for the options. If another screenshot exists with these
+	// options then it will be replaced.
+	Set(ctx context.Context, opts ScreenshotOptions, screenshot *Screenshot) error
+}
+
+// InmemCacheService is an in-mem screenshot cache.
+type InmemCacheService struct {
+	cache       *gocache.Cache
+	cacheHits   prometheus.Counter
+	cacheMisses prometheus.Counter
+}
+
+func NewInmemCacheService(expiration time.Duration, r prometheus.Registerer) CacheService {
+	return &InmemCacheService{
+		cache: gocache.New(expiration, time.Minute),
+		cacheHits: promauto.With(r).NewCounter(prometheus.CounterOpts{
+			Name:      "cache_hits_total",
+			Namespace: namespace,
+			Subsystem: subsystem,
+		}),
+		cacheMisses: promauto.With(r).NewCounter(prometheus.CounterOpts{
+			Name:      "cache_misses_total",
+			Namespace: namespace,
+			Subsystem: subsystem,
+		}),
+	}
+}
+
+func (s *InmemCacheService) Get(_ context.Context, opts ScreenshotOptions) (*Screenshot, bool) {
+	k := fmt.Sprintf("%s-%d-%s", opts.DashboardUID, opts.PanelID, opts.Theme)
+	if v, ok := s.cache.Get(k); ok {
+		defer s.cacheHits.Inc()
+		return v.(*Screenshot), true
+	}
+	defer s.cacheMisses.Inc()
+	return nil, false
+}
+
+func (s *InmemCacheService) Set(_ context.Context, opts ScreenshotOptions, screenshot *Screenshot) error {
+	k := fmt.Sprintf("%s-%d-%s", opts.DashboardUID, opts.PanelID, opts.Theme)
+	s.cache.Set(k, screenshot, 0)
+	return nil
+}
+
+type NoOpCacheService struct{}
+
+func (s *NoOpCacheService) Get(_ context.Context, _ ScreenshotOptions) (*Screenshot, bool) {
+	return nil, false
+}
+
+func (s *NoOpCacheService) Set(_ context.Context, _ ScreenshotOptions, _ *Screenshot) error {
+	return nil
+}

pkg/services/screenshot/cache_mock.go

@@ -0,0 +1,64 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: github.com/grafana/grafana/pkg/services/screenshot (interfaces: CacheService)
+
+// Package screenshot is a generated GoMock package.
+package screenshot
+
+import (
+	context "context"
+	reflect "reflect"
+
+	gomock "github.com/golang/mock/gomock"
+)
+
+// MockCacheService is a mock of CacheService interface.
+type MockCacheService struct {
+	ctrl     *gomock.Controller
+	recorder *MockCacheServiceMockRecorder
+}
+
+// MockCacheServiceMockRecorder is the mock recorder for MockCacheService.
+type MockCacheServiceMockRecorder struct {
+	mock *MockCacheService
+}
+
+// NewMockCacheService creates a new mock instance.
+func NewMockCacheService(ctrl *gomock.Controller) *MockCacheService {
+	mock := &MockCacheService{ctrl: ctrl}
+	mock.recorder = &MockCacheServiceMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use.
+func (m *MockCacheService) EXPECT() *MockCacheServiceMockRecorder {
+	return m.recorder
+}
+
+// Get mocks base method.
+func (m *MockCacheService) Get(arg0 context.Context, arg1 ScreenshotOptions) (*Screenshot, bool) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Get", arg0, arg1)
+	ret0, _ := ret[0].(*Screenshot)
+	ret1, _ := ret[1].(bool)
+	return ret0, ret1
+}
+
+// Get indicates an expected call of Get.
+func (mr *MockCacheServiceMockRecorder) Get(arg0, arg1 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Get", reflect.TypeOf((*MockCacheService)(nil).Get), arg0, arg1)
+}
+
+// Set mocks base method.
+func (m *MockCacheService) Set(arg0 context.Context, arg1 ScreenshotOptions, arg2 *Screenshot) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Set", arg0, arg1, arg2)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// Set indicates an expected call of Set.
+func (mr *MockCacheServiceMockRecorder) Set(arg0, arg1, arg2 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Set", reflect.TypeOf((*MockCacheService)(nil).Set), arg0, arg1, arg2)
+}

pkg/services/screenshot/cache_test.go

@@ -0,0 +1,37 @@
+package screenshot
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestInmemCacheService(t *testing.T) {
+	s := NewInmemCacheService(time.Second, prometheus.DefaultRegisterer)
+	ctx := context.Background()
+	opts := ScreenshotOptions{DashboardUID: "foo", PanelID: 1}
+
+	// should be a miss
+	actual, ok := s.Get(ctx, opts)
+	assert.False(t, ok)
+	assert.Nil(t, actual)
+
+	// should be a hit
+	expected := Screenshot{Path: "panel.png"}
+	require.NoError(t, s.Set(ctx, opts, &expected))
+	actual, ok = s.Get(ctx, opts)
+	assert.True(t, ok)
+	assert.Equal(t, expected, *actual)
+
+	// wait 1s and the cached screenshot should have expired
+	<-time.After(time.Second)
+
+	// should be a miss
+	actual, ok = s.Get(ctx, opts)
+	assert.False(t, ok)
+	assert.Nil(t, actual)
+}

pkg/services/screenshot/option.go

@@ -0,0 +1,43 @@
+package screenshot
+
+import (
+	"time"
+
+	"github.com/grafana/grafana/pkg/models"
+)
+
+var (
+	DefaultTheme   = models.ThemeDark
+	DefaultTimeout = 15 * time.Second
+	DefaultHeight  = 500
+	DefaultWidth   = 1000
+)
+
+// ScreenshotOptions are the options for taking a screenshot.
+type ScreenshotOptions struct {
+	DashboardUID string
+	PanelID      int64
+	Width        int
+	Height       int
+	Theme        models.Theme
+	Timeout      time.Duration
+}
+
+// SetDefaults sets default values for missing or invalid options.
+func (s ScreenshotOptions) SetDefaults() ScreenshotOptions {
+	if s.Width <= 0 {
+		s.Width = DefaultWidth
+	}
+	if s.Height <= 0 {
+		s.Height = DefaultHeight
+	}
+	switch s.Theme {
+	case models.ThemeDark, models.ThemeLight:
+	default:
+		s.Theme = DefaultTheme
+	}
+	if s.Timeout <= 0 {
+		s.Timeout = DefaultTimeout
+	}
+	return s
+}

pkg/services/screenshot/option_test.go

@@ -0,0 +1,55 @@
+package screenshot
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestScreenshotOptions(t *testing.T) {
+	o := ScreenshotOptions{}
+	assert.Equal(t, ScreenshotOptions{}, o)
+
+	o = o.SetDefaults()
+	assert.Equal(t, ScreenshotOptions{
+		Width:   DefaultWidth,
+		Height:  DefaultHeight,
+		Theme:   DefaultTheme,
+		Timeout: DefaultTimeout,
+	}, o)
+
+	o.Width = 100
+	o = o.SetDefaults()
+	assert.Equal(t, ScreenshotOptions{
+		Width:   100,
+		Height:  DefaultHeight,
+		Theme:   DefaultTheme,
+		Timeout: DefaultTimeout,
+	}, o)
+
+	o.Height = 100
+	o = o.SetDefaults()
+	assert.Equal(t, ScreenshotOptions{
+		Width:   100,
+		Height:  100,
+		Theme:   DefaultTheme,
+		Timeout: DefaultTimeout,
+	}, o)
+
+	o.Theme = "Not a theme"
+	o = o.SetDefaults()
+	assert.Equal(t, ScreenshotOptions{
+		Width:   100,
+		Height:  100,
+		Theme:   DefaultTheme,
+		Timeout: DefaultTimeout,
+	}, o)
+
+	o.Timeout = DefaultTimeout + 1
+	assert.Equal(t, ScreenshotOptions{
+		Width:   100,
+		Height:  100,
+		Theme:   DefaultTheme,
+		Timeout: DefaultTimeout + 1,
+	}, o)
+}

pkg/services/screenshot/ratelimit.go

@@ -0,0 +1,46 @@
+package screenshot
+
+import (
+	"context"
+)
+
+// A rate limiter restricts the number of screenshots that can be taken in parallel.
+//
+//go:generate mockgen -destination=ratelimit_mock.go -package=screenshot github.com/grafana/grafana/pkg/services/screenshot RateLimiter
+type RateLimiter interface {
+	// Do restricts the rate at which screenshots can be taken in parallel via screenshotFunc.
+	// It returns the result of screenshotFunc, or an error if either the context deadline
+	// has been exceeded or the context has been canceled while waiting its turn to call
+	// screenshotFunc.
+	Do(ctx context.Context, opts ScreenshotOptions, fn screenshotFunc) (*Screenshot, error)
+}
+
+// TokenRateLimiter is a rate limiter that uses a token bucket of fixed size N.
+type TokenRateLimiter struct {
+	tokens chan struct{}
+}
+
+func NewTokenRateLimiter(n int64) RateLimiter {
+	return &TokenRateLimiter{
+		tokens: make(chan struct{}, n),
+	}
+}
+
+func (s *TokenRateLimiter) Do(ctx context.Context, opts ScreenshotOptions, fn screenshotFunc) (*Screenshot, error) {
+	select {
+	// the context is canceled
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	// there is a token available
+	case s.tokens <- struct{}{}:
+		defer func() { <-s.tokens }()
+		return fn(ctx, opts)
+	}
+}
+
+// NoOpRateLimiter is a no-op rate limiter that has no limits.
+type NoOpRateLimiter struct{}
+
+func (b *NoOpRateLimiter) Do(ctx context.Context, opts ScreenshotOptions, fn screenshotFunc) (*Screenshot, error) {
+	return fn(ctx, opts)
+}

pkg/services/screenshot/ratelimit_mock.go

@@ -0,0 +1,50 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: github.com/grafana/grafana/pkg/services/screenshot (interfaces: RateLimiter)
+
+// Package screenshot is a generated GoMock package.
+package screenshot
+
+import (
+	context "context"
+	reflect "reflect"
+
+	gomock "github.com/golang/mock/gomock"
+)
+
+// MockRateLimiter is a mock of RateLimiter interface.
+type MockRateLimiter struct {
+	ctrl     *gomock.Controller
+	recorder *MockRateLimiterMockRecorder
+}
+
+// MockRateLimiterMockRecorder is the mock recorder for MockRateLimiter.
+type MockRateLimiterMockRecorder struct {
+	mock *MockRateLimiter
+}
+
+// NewMockRateLimiter creates a new mock instance.
+func NewMockRateLimiter(ctrl *gomock.Controller) *MockRateLimiter {
+	mock := &MockRateLimiter{ctrl: ctrl}
+	mock.recorder = &MockRateLimiterMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use.
+func (m *MockRateLimiter) EXPECT() *MockRateLimiterMockRecorder {
+	return m.recorder
+}
+
+// Do mocks base method.
+func (m *MockRateLimiter) Do(arg0 context.Context, arg1 ScreenshotOptions, arg2 screenshotFunc) (*Screenshot, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Do", arg0, arg1, arg2)
+	ret0, _ := ret[0].(*Screenshot)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// Do indicates an expected call of Do.
+func (mr *MockRateLimiterMockRecorder) Do(arg0, arg1, arg2 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Do", reflect.TypeOf((*MockRateLimiter)(nil).Do), arg0, arg1, arg2)
+}

pkg/services/screenshot/ratelimit_test.go

@@ -0,0 +1,54 @@
+package screenshot
+
+import (
+	"context"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/golang/mock/gomock"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestTokenRateLimiter(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	r := NewTokenRateLimiter(1)
+
+	ctx, cancelFunc := context.WithTimeout(context.Background(), time.Second)
+	defer cancelFunc()
+
+	var (
+		v  int64
+		wg sync.WaitGroup
+	)
+
+	testScreenshotFunc := func(ctx context.Context, opts ScreenshotOptions) (*Screenshot, error) {
+		// v should be 1 to show that no other goroutines acquired the token
+		atomic.AddInt64(&v, 1)
+		assert.Equal(t, int64(1), atomic.LoadInt64(&v))
+
+		// interrupt so other goroutines can attempt to acquire the token
+		<-time.After(time.Microsecond)
+
+		// v should be 0
+		atomic.AddInt64(&v, -1)
+		assert.Equal(t, int64(0), atomic.LoadInt64(&v))
+
+		return &Screenshot{}, nil
+	}
+
+	for i := 0; i < 10; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			screenshot, err := r.Do(ctx, ScreenshotOptions{}, testScreenshotFunc)
+			require.NoError(t, err)
+			assert.NotNil(t, screenshot)
+		}()
+	}
+	wg.Wait()
+}

pkg/services/screenshot/screenshot.go

@@ -9,12 +9,9 @@ import (
 	"strconv"
 	"time"
 
-	gocache "github.com/patrickmn/go-cache"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
-	"golang.org/x/sync/singleflight"
 
-	"github.com/grafana/grafana/pkg/components/imguploader"
 	"github.com/grafana/grafana/pkg/models"
 	"github.com/grafana/grafana/pkg/services/dashboards"
 	"github.com/grafana/grafana/pkg/services/org"
@@ -27,127 +24,38 @@ const (
 	subsystem = "screenshot"
 )
 
-var (
-	DefaultTheme   = models.ThemeDark
-	DefaultTimeout = 15 * time.Second
-	DefaultHeight  = 500
-	DefaultWidth   = 1000
-)
-
 var (
 	ErrScreenshotsUnavailable = errors.New("screenshots unavailable")
 )
 
-// Screenshot represents a screenshot of a dashboard in Grafana.
-//
-// A screenshot can have a Path and an URL if the screenshot is stored on disk
-// and uploaded to a cloud storage service or made accessible via the Grafana
-// HTTP server.
+// Screenshot represents a path to a screenshot on disk.
 type Screenshot struct {
 	Path string
-	URL  string
 }
 
-// ScreenshotOptions are the options for taking a screenshot.
-type ScreenshotOptions struct {
-	DashboardUID string
-	PanelID      int64
-	Width        int
-	Height       int
-	Theme        models.Theme
-	Timeout      time.Duration
-}
-
-// SetDefaults sets default values for missing or invalid options.
-func (s ScreenshotOptions) SetDefaults() ScreenshotOptions {
-	if s.Width <= 0 {
-		s.Width = DefaultWidth
-	}
-	if s.Height <= 0 {
-		s.Height = DefaultHeight
-	}
-	switch s.Theme {
-	case models.ThemeDark, models.ThemeLight:
-	default:
-		s.Theme = DefaultTheme
-	}
-	if s.Timeout <= 0 {
-		s.Timeout = DefaultTimeout
-	}
-	return s
-}
+type screenshotFunc func(ctx context.Context, opts ScreenshotOptions) (*Screenshot, error)
 
 // ScreenshotService is an interface for taking screenshots.
 //
-//go:generate mockgen -destination=mock.go -package=screenshot github.com/grafana/grafana/pkg/services/screenshot ScreenshotService
+//go:generate mockgen -destination=screenshot_mock.go -package=screenshot github.com/grafana/grafana/pkg/services/screenshot ScreenshotService
 type ScreenshotService interface {
 	Take(ctx context.Context, opts ScreenshotOptions) (*Screenshot, error)
 }
 
-// CachableScreenshotService caches screenshots.
-type CachableScreenshotService struct {
-	cache       *gocache.Cache
-	service     ScreenshotService
-	cacheHits   prometheus.Counter
-	cacheMisses prometheus.Counter
-}
+// HeadlessScreenshotService takes screenshots using a headless browser.
+type HeadlessScreenshotService struct {
+	ds dashboards.DashboardService
+	rs rendering.Service
 
-func NewCachableScreenshotService(r prometheus.Registerer, expiration time.Duration, service ScreenshotService) ScreenshotService {
-	return &CachableScreenshotService{
-		cache:   gocache.New(expiration, time.Minute),
-		service: service,
-		cacheHits: promauto.With(r).NewCounter(prometheus.CounterOpts{
-			Name:      "cache_hits_total",
-			Namespace: namespace,
-			Subsystem: subsystem,
-		}),
-		cacheMisses: promauto.With(r).NewCounter(prometheus.CounterOpts{
-			Name:      "cache_misses_total",
-			Namespace: namespace,
-			Subsystem: subsystem,
-		}),
-	}
-}
-
-// Take returns the screenshot from the cache or asks the service to take a
-// new screenshot and cache it before returning it.
-func (s *CachableScreenshotService) Take(ctx context.Context, opts ScreenshotOptions) (*Screenshot, error) {
-	k := fmt.Sprintf("%s-%d-%s", opts.DashboardUID, opts.PanelID, opts.Theme)
-
-	if v, ok := s.cache.Get(k); ok {
-		defer s.cacheHits.Inc()
-		return v.(*Screenshot), nil
-	}
-
-	defer s.cacheMisses.Inc()
-	screenshot, err := s.service.Take(ctx, opts)
-	if err != nil {
-		return nil, err
-	}
-
-	s.cache.Set(k, screenshot, 0)
-
-	return screenshot, nil
-}
-
-// NoopScreenshotService is a service that takes no-op screenshots.
-type NoopScreenshotService struct{}
-
-func (s *NoopScreenshotService) Take(_ context.Context, _ ScreenshotOptions) (*Screenshot, error) {
-	return &Screenshot{}, nil
-}
-
-// ObservableScreenshotService is a service that records metrics about screenshots.
-type ObservableScreenshotService struct {
-	service   ScreenshotService
 	duration  prometheus.Histogram
 	failures  *prometheus.CounterVec
 	successes prometheus.Counter
 }
 
-func NewObservableScreenshotService(r prometheus.Registerer, service ScreenshotService) ScreenshotService {
-	return &ObservableScreenshotService{
-		service: service,
+func NewHeadlessScreenshotService(ds dashboards.DashboardService, rs rendering.Service, r prometheus.Registerer) ScreenshotService {
+	return &HeadlessScreenshotService{
+		ds: ds,
+		rs: rs,
 		duration: promauto.With(r).NewHistogram(prometheus.HistogramOpts{
 			Name:      "duration_seconds",
 			Buckets:   []float64{0.1, 0.25, 0.5, 1, 2, 5, 10, 15},
@@ -167,53 +75,20 @@ func NewObservableScreenshotService(r prometheus.Registerer, service ScreenshotS
 	}
 }
 
-func (s *ObservableScreenshotService) Take(ctx context.Context, opts ScreenshotOptions) (*Screenshot, error) {
+// Take returns a screenshot of the panel. It returns an error if either the panel,
+// or the dashboard that contains the panel, does not exist, or the screenshot could not be
+// taken due to an error. It uses both the context and the timeout in ScreenshotOptions,
+// however the same context is used for both database queries and the request to the
+// rendering service, while the timeout in ScreenshotOptions is passed to the rendering service
+// where it is used as a client timeout. It is not recommended to pass a context without a deadline
+// and the context deadline should be at least as long as the timeout in ScreenshotOptions.
+func (s *HeadlessScreenshotService) Take(ctx context.Context, opts ScreenshotOptions) (*Screenshot, error) {
 	start := time.Now()
 	defer func() { s.duration.Observe(time.Since(start).Seconds()) }()
 
-	screenshot, err := s.service.Take(ctx, opts)
-	if err != nil {
-		if errors.Is(err, dashboards.ErrDashboardNotFound) {
-			defer s.failures.With(prometheus.Labels{
-				"reason": "dashboard_not_found",
-			}).Inc()
-		} else if errors.Is(err, context.Canceled) {
-			defer s.failures.With(prometheus.Labels{
-				"reason": "context_canceled",
-			}).Inc()
-		} else {
-			defer s.failures.With(prometheus.Labels{
-				"reason": "error",
-			}).Inc()
-		}
-	} else {
-		defer s.successes.Inc()
-	}
-	return screenshot, err
-}
-
-// RemoteRenderScreenshotService takes screenshots using a remote render service.
-type RemoteRenderScreenshotService struct {
-	ds dashboards.DashboardService
-	rs rendering.Service
-}
-
-func NewRemoteRenderScreenshotService(ds dashboards.DashboardService, rs rendering.Service) ScreenshotService {
-	return &RemoteRenderScreenshotService{
-		ds: ds,
-		rs: rs,
-	}
-}
-
-// Take returns a screenshot or an error if either the dashboard does not exist or
-// the service failed to screenshot the dashboard. It uses both the context and the
-// timeout in ScreenshotOptions, however the context is used in any database queries
-// and the request to the remote render service, while the timeout in ScreenshotOptions
-// is passed to the remote render service where it is used as a client timeout. It is
-// not recommended to pass a context without a deadline.
-func (s *RemoteRenderScreenshotService) Take(ctx context.Context, opts ScreenshotOptions) (*Screenshot, error) {
 	q := models.GetDashboardQuery{Uid: opts.DashboardUID}
 	if err := s.ds.GetDashboard(ctx, &q); err != nil {
+		s.instrumentError(err)
 		return nil, err
 	}
 
@@ -246,120 +121,41 @@ func (s *RemoteRenderScreenshotService) Take(ctx context.Context, opts Screensho
 
 	result, err := s.rs.Render(ctx, renderOpts, nil)
 	if err != nil {
+		s.instrumentError(err)
 		return nil, fmt.Errorf("failed to take screenshot: %w", err)
 	}
 
+	defer s.successes.Inc()
 	screenshot := Screenshot{Path: result.FilePath}
 	return &screenshot, nil
 }
 
+func (s *HeadlessScreenshotService) instrumentError(err error) {
+	if errors.Is(err, dashboards.ErrDashboardNotFound) {
+		defer s.failures.With(prometheus.Labels{
+			"reason": "dashboard_not_found",
+		}).Inc()
+	} else if errors.Is(err, context.Canceled) {
+		defer s.failures.With(prometheus.Labels{
+			"reason": "context_canceled",
+		}).Inc()
+	} else {
+		defer s.failures.With(prometheus.Labels{
+			"reason": "error",
+		}).Inc()
+	}
+}
+
+// NoOpScreenshotService is a service that takes no-op screenshots.
+type NoOpScreenshotService struct{}
+
+func (s *NoOpScreenshotService) Take(_ context.Context, _ ScreenshotOptions) (*Screenshot, error) {
+	return &Screenshot{}, nil
+}
+
+// ScreenshotUnavailableService is a service that returns ErrScreenshotsUnavailable.
 type ScreenshotUnavailableService struct{}
 
 func (s *ScreenshotUnavailableService) Take(_ context.Context, _ ScreenshotOptions) (*Screenshot, error) {
 	return nil, ErrScreenshotsUnavailable
 }
-
-// SingleFlightScreenshotService prevents duplicate screenshots.
-type SingleFlightScreenshotService struct {
-	f       singleflight.Group
-	service ScreenshotService
-}
-
-func NewSingleFlightScreenshotService(service ScreenshotService) ScreenshotService {
-	return &SingleFlightScreenshotService{service: service}
-}
-
-// Take returns a screenshot or an error. It ensures that at most one screenshot
-// can be taken at a time for the same dashboard and theme. Duplicate screenshots
-// wait for the first screenshot to complete and receive the same screenshot.
-func (s *SingleFlightScreenshotService) Take(ctx context.Context, opts ScreenshotOptions) (*Screenshot, error) {
-	k := fmt.Sprintf("%s-%d-%s", opts.DashboardUID, opts.PanelID, opts.Theme)
-
-	v, err, _ := s.f.Do(k, func() (interface{}, error) {
-		return s.service.Take(ctx, opts)
-	})
-	if err != nil {
-		return nil, err
-	}
-
-	screenshot := v.(*Screenshot)
-	return screenshot, err
-}
-
-// RateLimitScreenshotService ensures that at most N screenshots can be taken
-// at a time.
-type RateLimitScreenshotService struct {
-	service ScreenshotService
-	tokens  chan struct{}
-}
-
-func NewRateLimitScreenshotService(service ScreenshotService, n int64) ScreenshotService {
-	return &RateLimitScreenshotService{
-		service: service,
-		tokens:  make(chan struct{}, n),
-	}
-}
-
-// Take returns a screenshot or an error. It ensures that at most N screenshots
-// can be taken at a time. The service has N tokens such that a token is consumed
-// at the start of a screenshot and returned when the screenshot has either
-// succeeded or failed. A screenshot can timeout if the context is canceled
-// while waiting for a token or while the screenshot is being taken.
-func (s *RateLimitScreenshotService) Take(ctx context.Context, opts ScreenshotOptions) (*Screenshot, error) {
-	select {
-	// the context is canceled
-	case <-ctx.Done():
-		return nil, ctx.Err()
-	// there is a token available
-	case s.tokens <- struct{}{}:
-	}
-	// acquired token must be returned
-	defer func() {
-		<-s.tokens
-	}()
-	return s.service.Take(ctx, opts)
-}
-
-// UploadingScreenshotService uploads taken screenshots.
-type UploadingScreenshotService struct {
-	service         ScreenshotService
-	uploader        imguploader.ImageUploader
-	uploadFailures  prometheus.Counter
-	uploadSuccesses prometheus.Counter
-}
-
-func NewUploadingScreenshotService(r prometheus.Registerer, service ScreenshotService, uploader imguploader.ImageUploader) ScreenshotService {
-	return &UploadingScreenshotService{
-		service:  service,
-		uploader: uploader,
-		uploadFailures: promauto.With(r).NewCounter(prometheus.CounterOpts{
-			Name:      "upload_failures_total",
-			Namespace: namespace,
-			Subsystem: subsystem,
-		}),
-		uploadSuccesses: promauto.With(r).NewCounter(prometheus.CounterOpts{
-			Name:      "upload_successes_total",
-			Namespace: namespace,
-			Subsystem: subsystem,
-		}),
-	}
-}
-
-// Take uploads a screenshot with a path and returns a new screenshot with the
-// unmodified path and a URL. It returns the unmodified screenshot on error.
-func (s *UploadingScreenshotService) Take(ctx context.Context, opts ScreenshotOptions) (*Screenshot, error) {
-	screenshot, err := s.service.Take(ctx, opts)
-	if err != nil {
-		return nil, err
-	}
-
-	url, err := s.uploader.Upload(ctx, screenshot.Path)
-	if err != nil {
-		defer s.uploadFailures.Inc()
-		return screenshot, fmt.Errorf("failed to upload screenshot: %w", err)
-	}
-	screenshot.URL = url
-
-	defer s.uploadSuccesses.Inc()
-	return screenshot, nil
-}

pkg/services/screenshot/screenshot_test.go

@@ -2,12 +2,8 @@ package screenshot
 
 import (
 	"context"
-	"errors"
 	"fmt"
-	"sync"
-	"sync/atomic"
 	"testing"
-	"time"
 
 	"github.com/golang/mock/gomock"
 	"github.com/prometheus/client_golang/prometheus"
@@ -15,7 +11,6 @@ import (
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 
-	"github.com/grafana/grafana/pkg/components/imguploader"
 	"github.com/grafana/grafana/pkg/models"
 	"github.com/grafana/grafana/pkg/services/dashboards"
 	"github.com/grafana/grafana/pkg/services/org"
@@ -23,61 +18,13 @@ import (
 	"github.com/grafana/grafana/pkg/setting"
 )
 
-func TestScreenshotOptions(t *testing.T) {
-	o := ScreenshotOptions{}
-	assert.Equal(t, ScreenshotOptions{}, o)
-
-	o = o.SetDefaults()
-	assert.Equal(t, ScreenshotOptions{
-		Width:   DefaultWidth,
-		Height:  DefaultHeight,
-		Theme:   DefaultTheme,
-		Timeout: DefaultTimeout,
-	}, o)
-
-	o.Width = 100
-	o = o.SetDefaults()
-	assert.Equal(t, ScreenshotOptions{
-		Width:   100,
-		Height:  DefaultHeight,
-		Theme:   DefaultTheme,
-		Timeout: DefaultTimeout,
-	}, o)
-
-	o.Height = 100
-	o = o.SetDefaults()
-	assert.Equal(t, ScreenshotOptions{
-		Width:   100,
-		Height:  100,
-		Theme:   DefaultTheme,
-		Timeout: DefaultTimeout,
-	}, o)
-
-	o.Theme = "Not a theme"
-	o = o.SetDefaults()
-	assert.Equal(t, ScreenshotOptions{
-		Width:   100,
-		Height:  100,
-		Theme:   DefaultTheme,
-		Timeout: DefaultTimeout,
-	}, o)
-
-	o.Timeout = DefaultTimeout + 1
-	assert.Equal(t, ScreenshotOptions{
-		Width:   100,
-		Height:  100,
-		Theme:   DefaultTheme,
-		Timeout: DefaultTimeout + 1,
-	}, o)
-}
-
-func TestBrowserScreenshotService(t *testing.T) {
+func TestHeadlessScreenshotService(t *testing.T) {
 	c := gomock.NewController(t)
 	defer c.Finish()
 
 	d := dashboards.FakeDashboardService{}
 	r := rendering.NewMockService(c)
-	s := NewRemoteRenderScreenshotService(&d, r)
+	s := NewHeadlessScreenshotService(&d, r, prometheus.NewRegistry())
 
 	// a non-existent dashboard should return error
 	d.On("GetDashboard", mock.Anything, mock.AnythingOfType("*models.GetDashboardQuery")).Return(dashboards.ErrDashboardNotFound).Once()
@@ -129,39 +76,8 @@ func TestBrowserScreenshotService(t *testing.T) {
 	assert.Nil(t, screenshot)
 }
 
-func TestCachableScreenshotService(t *testing.T) {
-	c := gomock.NewController(t)
-	defer c.Finish()
-
-	m := NewMockScreenshotService(c)
-	s := NewCachableScreenshotService(prometheus.DefaultRegisterer, time.Second, m)
-
-	ctx := context.Background()
-	opts := ScreenshotOptions{DashboardUID: "foo", PanelID: 1}
-
-	// should be a miss and ask the mock service to take a screenshot
-	m.EXPECT().Take(ctx, opts).Return(&Screenshot{Path: "panel.png"}, nil)
-	screenshot, err := s.Take(ctx, opts)
-	require.NoError(t, err)
-	assert.Equal(t, Screenshot{Path: "panel.png"}, *screenshot)
-
-	// should be a hit
-	screenshot, err = s.Take(ctx, opts)
-	require.NoError(t, err)
-	assert.Equal(t, Screenshot{Path: "panel.png"}, *screenshot)
-
-	// wait 1s and the cached screenshot should have expired
-	<-time.After(time.Second)
-
-	// should be a miss as the cached screenshot has expired
-	m.EXPECT().Take(ctx, opts).Return(&Screenshot{Path: "panel.png"}, nil)
-	screenshot, err = s.Take(ctx, opts)
-	require.NoError(t, err)
-	assert.Equal(t, Screenshot{Path: "panel.png"}, *screenshot)
-}
-
-func TestNoopScreenshotService(t *testing.T) {
-	s := NoopScreenshotService{}
+func TestNoOpScreenshotService(t *testing.T) {
+	s := NoOpScreenshotService{}
 	screenshot, err := s.Take(context.Background(), ScreenshotOptions{})
 	assert.NoError(t, err)
 	assert.NotNil(t, screenshot)
@@ -173,166 +89,3 @@ func TestScreenshotUnavailableService(t *testing.T) {
 	assert.Equal(t, err, ErrScreenshotsUnavailable)
 	assert.Nil(t, screenshot)
 }
-
-func TestSingleFlightScreenshotService(t *testing.T) {
-	c := gomock.NewController(t)
-	defer c.Finish()
-
-	m := NewMockScreenshotService(c)
-	s := NewSingleFlightScreenshotService(m)
-
-	ctx := context.Background()
-	opts := ScreenshotOptions{DashboardUID: "foo", PanelID: 1}
-
-	// expect 1 invocation of the mock service for the same options
-	m.EXPECT().Take(ctx, opts).
-		Do(func(_ context.Context, _ ScreenshotOptions) { <-time.After(time.Second) }).
-		Return(&Screenshot{Path: "panel.png"}, nil)
-
-	var wg sync.WaitGroup
-	for i := 0; i < 5; i++ {
-		wg.Add(1)
-		go func() {
-			defer wg.Done()
-			screenshot, err := s.Take(ctx, opts)
-			require.NoError(t, err)
-			assert.Equal(t, Screenshot{Path: "panel.png"}, *screenshot)
-		}()
-	}
-	wg.Wait()
-
-	// expect two invocations of the mock service for different dashboards
-	opts1 := ScreenshotOptions{DashboardUID: "foo", PanelID: 1}
-	opts2 := ScreenshotOptions{DashboardUID: "bar", PanelID: 1}
-	m.EXPECT().Take(ctx, opts1).Return(&Screenshot{Path: "foo.png"}, nil)
-	m.EXPECT().Take(ctx, opts2).Return(&Screenshot{Path: "bar.png"}, nil)
-
-	wg.Add(2)
-	go func() {
-		defer wg.Done()
-		screenshot, err := s.Take(ctx, opts1)
-		require.NoError(t, err)
-		assert.Equal(t, Screenshot{Path: "foo.png"}, *screenshot)
-	}()
-	go func() {
-		defer wg.Done()
-		screenshot, err := s.Take(ctx, opts2)
-		require.NoError(t, err)
-		assert.Equal(t, Screenshot{Path: "bar.png"}, *screenshot)
-	}()
-	wg.Wait()
-
-	// expect two invocations of the mock service for different panels in the same dashboard
-	opts1 = ScreenshotOptions{DashboardUID: "foo", PanelID: 1}
-	opts2 = ScreenshotOptions{DashboardUID: "foo", PanelID: 2}
-	m.EXPECT().Take(ctx, opts1).Return(&Screenshot{Path: "panel1.png"}, nil)
-	m.EXPECT().Take(ctx, opts2).Return(&Screenshot{Path: "panel2.png"}, nil)
-
-	wg.Add(2)
-	go func() {
-		defer wg.Done()
-		screenshot, err := s.Take(ctx, opts1)
-		require.NoError(t, err)
-		assert.Equal(t, Screenshot{Path: "panel1.png"}, *screenshot)
-	}()
-	go func() {
-		defer wg.Done()
-		screenshot, err := s.Take(ctx, opts2)
-		require.NoError(t, err)
-		assert.Equal(t, Screenshot{Path: "panel2.png"}, *screenshot)
-	}()
-	wg.Wait()
-
-	// expect two invocations of the mock service for different panels in the same dashboard
-	opts1 = ScreenshotOptions{DashboardUID: "foo", PanelID: 1, Theme: models.ThemeDark}
-	opts2 = ScreenshotOptions{DashboardUID: "foo", PanelID: 1, Theme: models.ThemeLight}
-	m.EXPECT().Take(ctx, opts1).Return(&Screenshot{Path: "dark.png"}, nil)
-	m.EXPECT().Take(ctx, opts2).Return(&Screenshot{Path: "light.png"}, nil)
-
-	wg.Add(2)
-	go func() {
-		defer wg.Done()
-		screenshot, err := s.Take(ctx, opts1)
-		require.NoError(t, err)
-		assert.Equal(t, Screenshot{Path: "dark.png"}, *screenshot)
-	}()
-	go func() {
-		defer wg.Done()
-		screenshot, err := s.Take(ctx, opts2)
-		require.NoError(t, err)
-		assert.Equal(t, Screenshot{Path: "light.png"}, *screenshot)
-	}()
-	wg.Wait()
-}
-
-func TestRateLimitScreenshotService(t *testing.T) {
-	c := gomock.NewController(t)
-	defer c.Finish()
-
-	m := NewMockScreenshotService(c)
-	s := NewRateLimitScreenshotService(m, 1)
-	ctx := context.Background()
-	opts := ScreenshotOptions{DashboardUID: "foo", PanelID: 1}
-
-	var v int64
-	for i := 0; i < 10; i++ {
-		m.EXPECT().Take(ctx, opts).
-			Do(func(_ context.Context, _ ScreenshotOptions) {
-				// v should be 0 to show that no tokens have been acquired
-				assert.Equal(t, int64(0), atomic.LoadInt64(&v))
-				atomic.AddInt64(&v, 1)
-				assert.Equal(t, int64(1), atomic.LoadInt64(&v))
-
-				// interrupt so other goroutines can attempt to acquire the token
-				<-time.After(time.Microsecond)
-
-				// v should be 1 to show that no other goroutines acquired the token
-				assert.Equal(t, int64(1), atomic.LoadInt64(&v))
-				atomic.AddInt64(&v, -1)
-				assert.Equal(t, int64(0), atomic.LoadInt64(&v))
-			}).
-			Return(&Screenshot{Path: "foo.png"}, nil)
-	}
-
-	var wg sync.WaitGroup
-	for i := 0; i < 10; i++ {
-		wg.Add(1)
-		go func() {
-			defer wg.Done()
-			result, err := s.Take(ctx, opts)
-			require.NoError(t, err)
-			assert.Equal(t, Screenshot{Path: "foo.png"}, *result)
-		}()
-	}
-	wg.Wait()
-}
-
-func TestUploadingScreenshotService(t *testing.T) {
-	c := gomock.NewController(t)
-	defer c.Finish()
-
-	m := NewMockScreenshotService(c)
-	u := imguploader.NewMockImageUploader(c)
-	s := NewUploadingScreenshotService(prometheus.DefaultRegisterer, m, u)
-
-	ctx := context.Background()
-	opts := ScreenshotOptions{DashboardUID: "foo", PanelID: 1}
-
-	m.EXPECT().Take(ctx, opts).Return(&Screenshot{Path: "foo.png"}, nil)
-	u.EXPECT().Upload(ctx, "foo.png").Return("https://example.com/foo.png", nil)
-	screenshot, err := s.Take(ctx, opts)
-	require.NoError(t, err)
-	assert.Equal(t, Screenshot{
-		Path: "foo.png",
-		URL:  "https://example.com/foo.png",
-	}, *screenshot)
-
-	// error on upload should still return screenshot on disk
-	m.EXPECT().Take(ctx, opts).Return(&Screenshot{Path: "foo.png"}, nil)
-	u.EXPECT().Upload(ctx, "foo.png").Return("", errors.New("service is unavailable"))
-	screenshot, err = s.Take(ctx, opts)
-	assert.EqualError(t, err, "failed to upload screenshot: service is unavailable")
-	assert.Equal(t, Screenshot{
-		Path: "foo.png",
-	}, *screenshot)
-}