Support retry

DaoCloud · May 10, 2024 · 25ad7eb · 25ad7eb
1 parent bb5ea7b
commit 25ad7eb
Show file tree

Hide file tree

Showing 4 changed files with 37 additions and 1 deletion.
diff --git 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
@@ -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
@@ -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
@@ -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=