diff --git a/cmd/crproxy/main.go b/cmd/crproxy/main.go index 9451e1a..efce778 100644 --- a/cmd/crproxy/main.go +++ b/cmd/crproxy/main.go @@ -9,6 +9,7 @@ import ( "os" "slices" "strings" + "time" "github.com/gorilla/handlers" "github.com/spf13/pflag" @@ -24,6 +25,8 @@ var ( blobsSpeedLimit string totalBlobsSpeedLimit string blockImageList []string + retry int + retryInterval time.Duration ) func init() { @@ -33,6 +36,8 @@ func init() { pflag.StringVar(&blobsSpeedLimit, "blobs-speed-limit", "", "blobs speed limit per second (default unlimited)") pflag.StringVar(&totalBlobsSpeedLimit, "total-blobs-speed-limit", "", "total blobs speed limit per second (default unlimited)") pflag.StringSliceVar(&blockImageList, "block-image-list", nil, "block image list") + pflag.IntVar(&retry, "retry", 0, "retry times") + pflag.DurationVar(&retryInterval, "retry-interval", 0, "retry interval") pflag.Parse() } @@ -134,6 +139,10 @@ func main() { opts = append(opts, crproxy.WithTotalBlobsSpeedLimit(b)) } + if retry > 0 { + opts = append(opts, crproxy.WithRetry(retry, retryInterval)) + } + crp, err := crproxy.NewCRProxy(opts...) if err != nil { logger.Println("failed to NewCRProxy:", err) diff --git a/crproxy.go b/crproxy.go index 5234f8f..639e86b 100644 --- a/crproxy.go +++ b/crproxy.go @@ -14,6 +14,7 @@ import ( "github.com/distribution/distribution/v3/registry/client/auth/challenge" "github.com/distribution/distribution/v3/registry/client/transport" "github.com/wzshiming/geario" + "github.com/wzshiming/httpseek" "github.com/wzshiming/lru" ) @@ -44,6 +45,8 @@ type CRProxy struct { totalBlobsSpeedLimit *geario.Gear blobsSpeedLimit *geario.B blockFunc func(*PathInfo) bool + retry int + retryInterval time.Duration } type Option func(c *CRProxy) @@ -111,6 +114,13 @@ func WithBlockFunc(blockFunc func(info *PathInfo) bool) Option { } } +func WithRetry(retry int, retryInterval time.Duration) Option { + return func(c *CRProxy) { + c.retry = retry + c.retryInterval = retryInterval + } +} + func NewCRProxy(opts ...Option) (*CRProxy, error) { c := &CRProxy{ challengeManager: challenge.NewSimpleManager(), @@ -177,6 +187,19 @@ func (c *CRProxy) getClientset(host string, image string) *http.Client { } } + if c.retryInterval > 0 { + tr = httpseek.NewMustReaderTransport(tr, func(request *http.Request, retry int, err error) error { + if c.retry > 0 && retry >= c.retry { + return err + } + if c.logger != nil { + c.logger.Println("Retry", request.URL, retry, err) + } + time.Sleep(c.retryInterval) + return nil + }) + } + tr = transport.NewTransport(tr, auth.NewAuthorizer(c.challengeManager, authHandler)) client := &http.Client{ diff --git a/go.mod b/go.mod index e9215e5..ba5e6cc 100644 --- a/go.mod +++ b/go.mod @@ -1,12 +1,13 @@ module github.com/wzshiming/crproxy -go 1.19 +go 1.21 require ( github.com/distribution/distribution/v3 v3.0.0-20220907155224-78b9c98c5c31 github.com/gorilla/handlers v1.5.1 github.com/spf13/pflag v1.0.3 github.com/wzshiming/geario v0.0.0-20220512063849-ddc203ba7487 + github.com/wzshiming/httpseek v0.1.0 github.com/wzshiming/lru v0.1.0 ) diff --git a/go.sum b/go.sum index 9405477..318665a 100644 --- a/go.sum +++ b/go.sum @@ -59,6 +59,7 @@ github.com/distribution/distribution/v3 v3.0.0-20220907155224-78b9c98c5c31/go.mo github.com/docker/go-metrics v0.0.1 h1:AgB/0SvBxihN0X8OR4SjsblXkbMvalQ8cjmtKQ2rQV8= github.com/docker/go-metrics v0.0.1/go.mod h1:cG1hvH2utMXtqgqqYE9plW6lDxS3/5ayHzueweSI3Vw= github.com/docker/libtrust v0.0.0-20150114040149-fa567046d9b1 h1:ZClxb8laGDf5arXfYcAtECDFgAgHklGI8CxgjHnXKJ4= +github.com/docker/libtrust v0.0.0-20150114040149-fa567046d9b1/go.mod h1:cyGadeNEkKy96OOhEzfZl+yxihPEzKnqJwvfuSUqbZE= github.com/envoyproxy/go-control-plane v0.9.0/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4= github.com/envoyproxy/go-control-plane v0.9.1-0.20191026205805-5f8ba28d4473/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4= github.com/envoyproxy/go-control-plane v0.9.4/go.mod h1:6rpuAdCZL397s3pYoYcLgu1mIlRU8Am5FuJP05cCM98= @@ -219,6 +220,8 @@ github.com/stretchr/testify v1.8.0 h1:pSgiaMZlXftHpm5L7V1+rVB+AZJydKsMxsQBIJw4PK github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= github.com/wzshiming/geario v0.0.0-20220512063849-ddc203ba7487 h1:5YN2BY58L2/5N7olluOPqJmmybcio0+DFY6nwH5gc8c= github.com/wzshiming/geario v0.0.0-20220512063849-ddc203ba7487/go.mod h1:Fodw3HJvNUS+/MgqXCRp9iYLQfynAu/LKXGOWoX+D/Q= +github.com/wzshiming/httpseek v0.1.0 h1:lEgL7EBELT/VV9UaTp+m3kw5Pe1KOUdY+IPnKkag6tI= +github.com/wzshiming/httpseek v0.1.0/go.mod h1:YoZhlLIwNjTBDXIT8NpK5zRjOgZouRXPaBfjVXdqMMs= github.com/wzshiming/lru v0.1.0 h1:937SBBo9lDBRMivJqF46eVK2Q3omRWnN7hWpr2i3xZg= github.com/wzshiming/lru v0.1.0/go.mod h1:uaI2W/Gx3r1gHIc3FG38PKK1idvMXhgY8lToXaBsDIQ= github.com/yuin/goldmark v1.1.25/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=