grafana/pkg/services/screenshot/ratelimit.go
George Robinson bad4f7fec5
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.
2022-09-21 10:25:07 +01:00

47 lines
1.4 KiB
Go

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)
}