1berry/spegel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Use a single rate limiter for all requests.

Browse Source
This commit is contained in:
Anders Qvist 2024-02-21 18:15:31 +01:00
parent a5c16a73d2
commit bb68b64aa1
No known key found for this signature in database
GPG Key ID: 45617D5901E90925
3 changed files with 14 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
pkg/registry/registry.go
View File

@ -17,6 +17,7 @@ import (
"github.com/go-logr/logr"
"github.com/opencontainers/go-digest"
pkggin "github.com/xenitab/pkg/gin"
"golang.org/x/time/rate"
"github.com/xenitab/spegel/pkg/metrics"
"github.com/xenitab/spegel/pkg/oci"
@ -26,6 +27,7 @@ import (
const (
MirroredHeaderKey = "X-Spegel-Mirrored"
burstLimit = 1024 * 1024 * 1024 // 1GB
)
type Option func(*Registry)
@ -62,12 +64,14 @@ func WithLocalAddress(localAddr string) Option {
func WithBlobSpeed(blobSpeed throttle.Byterate) Option {
return func(r *Registry) {
r.blobSpeed = &blobSpeed
limiter := rate.NewLimiter(rate.Limit(blobSpeed), burstLimit)
limiter.AllowN(time.Now(), burstLimit)
r.rateLimiter = limiter
}
}
type Registry struct {
blobSpeed *throttle.Byterate
rateLimiter *rate.Limiter
ociClient oci.Client
router routing.Router
transport http.RoundTripper
@ -304,8 +308,8 @@ func (r *Registry) handleBlob(c *gin.Context, dgst digest.Digest) {
return
}
var writer io.Writer = c.Writer
if r.blobSpeed != nil {
writer = throttle.NewWriter(c.Writer, *r.blobSpeed)
if r.rateLimiter != nil {
writer = throttle.NewWriter(c.Writer, r.rateLimiter)
}
err = r.ociClient.CopyLayer(c, dgst, writer)
if err != nil {

6
pkg/throttle/writer.go
View File

@ -8,16 +8,12 @@ import (
"golang.org/x/time/rate"
)
const burstLimit = 1024 * 1024 * 1024 // 1GB
type writer struct {
limiter *rate.Limiter
writer io.Writer
}
func NewWriter(w io.Writer, br Byterate) io.Writer {
limiter := rate.NewLimiter(rate.Limit(br), burstLimit)
limiter.AllowN(time.Now(), burstLimit)
func NewWriter(w io.Writer, limiter *rate.Limiter) io.Writer {
return &writer{
limiter: limiter,
writer: w,

7
pkg/throttle/writer_test.go
View File

@ -6,11 +6,14 @@ import (
"time"
"github.com/stretchr/testify/require"
"golang.org/x/time/rate"
)
func TestWriter(t *testing.T) {
br := 500 * Bps
w := NewWriter(bytes.NewBuffer([]byte{}), br)
limit := rate.Limit(500 * Bps)
limiter := rate.NewLimiter(limit, 1024*1024)
limiter.AllowN(time.Now(), 1024*1024)
w := NewWriter(bytes.NewBuffer([]byte{}), limiter)
chunkSize := 100
start := time.Now()
for i := 0; i < 10; i++ {