feat: add rate limiter for alibaba cloud acr adapter

src/pkg/reg/adapter/tencentcr/ratelimiter.go → src/lib/rate_limiter_transport.go

@@ -12,11 +12,10 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package tencentcr
+package lib
 
 import (
 	"net/http"
-	"sync"
 
 	"go.uber.org/ratelimit"
 )
@@ -33,19 +32,9 @@ func (t limitTransport) RoundTrip(req *http.Request) (*http.Response, error) {
 	return t.RoundTripper.RoundTrip(req)
 }
 
-var limiterOnce sync.Once
-var limiter ratelimit.Limiter
-
-func newLimiter(rate int) ratelimit.Limiter {
-	limiterOnce.Do(func() {
-		limiter = ratelimit.New(rate)
-	})
-	return limiter
-}
-
-func newRateLimitedTransport(rate int, transport http.RoundTripper) http.RoundTripper {
+func NewRateLimitedTransport(rate int, transport http.RoundTripper) http.RoundTripper {
 	return &limitTransport{
 		RoundTripper: transport,
-		limiter:      newLimiter(rate),
+		limiter:      ratelimit.New(rate),
 	}
 }

src/lib/rate_limiter_transport_test.go

@@ -0,0 +1,45 @@
+// Copyright Project Harbor Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package lib
+
+import (
+	"net/http"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func Test_newRateLimitedTransport(t *testing.T) {
+	tests := []struct {
+		name      string
+		rate      int
+		transport http.RoundTripper
+	}{
+		{"1qps", 1, http.DefaultTransport},
+	}
+	var req = &http.Request{}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := NewRateLimitedTransport(tt.rate, tt.transport)
+			start := time.Now()
+			for i := 0; i <= tt.rate; i++ {
+				got.RoundTrip(req)
+			}
+			used := int64(time.Since(start).Milliseconds()) / int64(tt.rate)
+			assert.GreaterOrEqualf(t, used/int64(tt.rate), int64(1e3/tt.rate), "used %d ms per req", used/int64(tt.rate))
+		})
+	}
+}

src/pkg/reg/adapter/aliacr/adapter.go

@@ -17,12 +17,16 @@ package aliacr
 import (
 	"errors"
 	"fmt"
+	"net/http"
+	"os"
 	"path/filepath"
 	"regexp"
+	"strconv"
 	"strings"
 
 	commonhttp "github.com/goharbor/harbor/src/common/http"
 	"github.com/goharbor/harbor/src/common/utils"
+	"github.com/goharbor/harbor/src/lib"
 	"github.com/goharbor/harbor/src/lib/log"
 	adp "github.com/goharbor/harbor/src/pkg/reg/adapter"
 	"github.com/goharbor/harbor/src/pkg/reg/adapter/native"
@@ -32,7 +36,17 @@ import (
 	"github.com/goharbor/harbor/src/pkg/registry/auth/bearer"
 )
 
+var rateLimiterTransport http.RoundTripper
+
+const acrQPSLimit = 15
+
 func init() {
+	var envAcrQPSLimit, _ = strconv.Atoi(os.Getenv("REG_ADAPTER_ACR_QPS_LIMIT"))
+	if envAcrQPSLimit > acrQPSLimit || envAcrQPSLimit < 1 {
+		envAcrQPSLimit = acrQPSLimit
+	}
+	rateLimiterTransport = lib.NewRateLimitedTransport(envAcrQPSLimit, commonhttp.NewTransport())
+
 	if err := adp.RegisterFactory(model.RegistryTypeAliAcr, new(factory)); err != nil {
 		log.Errorf("failed to register factory for %s: %v", model.RegistryTypeAliAcr, err)
 		return
@@ -101,12 +115,12 @@ func newAdapter(registry *model.Registry) (*adapter, error) {
 
 	var acrAPI openapi
 	if !info.IsACREE {
-		acrAPI, err = newAcrOpenapi(registry.Credential.AccessKey, registry.Credential.AccessSecret, info.RegionID)
+		acrAPI, err = newAcrOpenapi(registry.Credential.AccessKey, registry.Credential.AccessSecret, info.RegionID, rateLimiterTransport)
 		if err != nil {
 			return nil, err
 		}
 	} else {
-		acrAPI, err = newAcreeOpenapi(registry.Credential.AccessKey, registry.Credential.AccessSecret, info.RegionID, info.InstanceID)
+		acrAPI, err = newAcreeOpenapi(registry.Credential.AccessKey, registry.Credential.AccessSecret, info.RegionID, info.InstanceID, rateLimiterTransport)
 		if err != nil {
 			return nil, err
 		}

src/pkg/reg/adapter/aliacr/openapi.go

@@ -17,12 +17,17 @@ package aliacr
 import (
 	"encoding/json"
 	"fmt"
+	"net/http"
 	"strconv"
 	"time"
 
+	"github.com/aliyun/alibaba-cloud-sdk-go/sdk"
+	"github.com/aliyun/alibaba-cloud-sdk-go/sdk/auth/credentials"
 	"github.com/aliyun/alibaba-cloud-sdk-go/sdk/requests"
 	"github.com/aliyun/alibaba-cloud-sdk-go/services/cr"
 	"github.com/aliyun/alibaba-cloud-sdk-go/services/cr_ee"
+
+	"github.com/goharbor/harbor/src/pkg/registry"
 )
 
 type repository struct {
@@ -57,8 +62,13 @@ type acrOpenapi struct {
 var _ openapi = &acrOpenapi{}
 
 // newAcrOpenapi creates a new acrOpenapi instance.
-func newAcrOpenapi(accessKeyID string, accessKeySecret string, regionID string) (openapi, error) {
-	client, err := cr.NewClientWithAccessKey(regionID, accessKeyID, accessKeySecret)
+func newAcrOpenapi(accessKeyID string, accessKeySecret string, regionID string, rateLimiter http.RoundTripper) (openapi, error) {
+	config := sdk.NewConfig().
+		WithUserAgent(registry.UserAgent).
+		WithAutoRetry(true)
+	config.Transport = rateLimiter
+
+	client, err := cr.NewClientWithOptions(regionID, config, credentials.NewAccessKeyCredential(accessKeyID, accessKeySecret))
 	if err != nil {
 		return nil, err
 	}
@@ -182,8 +192,13 @@ type acreeOpenapi struct {
 var _ openapi = &acreeOpenapi{}
 
 // newAcreeOpenapi creates a new acreeOpenapi instance.
-func newAcreeOpenapi(accessKeyID string, accessKeySecret string, regionID string, instanceID string) (openapi, error) {
-	client, err := cr_ee.NewClientWithAccessKey(regionID, accessKeyID, accessKeySecret)
+func newAcreeOpenapi(accessKeyID string, accessKeySecret string, regionID string, instanceID string, rateLimiter http.RoundTripper) (openapi, error) {
+	config := sdk.NewConfig().
+		WithUserAgent(registry.UserAgent).
+		WithAutoRetry(true)
+	config.Transport = rateLimiter
+
+	client, err := cr_ee.NewClientWithOptions(regionID, config, credentials.NewAccessKeyCredential(accessKeyID, accessKeySecret))
 	if err != nil {
 		return nil, err
 	}

src/pkg/reg/adapter/tencentcr/adapter.go

@@ -30,6 +30,7 @@ import (
 	tcr "github.com/tencentcloud/tencentcloud-sdk-go/tencentcloud/tcr/v20190924"
 
 	commonhttp "github.com/goharbor/harbor/src/common/http"
+	"github.com/goharbor/harbor/src/lib"
 	"github.com/goharbor/harbor/src/lib/log"
 	adp "github.com/goharbor/harbor/src/pkg/reg/adapter"
 	"github.com/goharbor/harbor/src/pkg/reg/adapter/native"
@@ -40,13 +41,15 @@ import (
 
 var (
 	errInvalidTcrEndpoint = errors.New("[tencent-tcr.newAdapter] Invalid TCR instance endpoint")
+	rateLimiterTransport  http.RoundTripper
 )
 
 func init() {
-	var envTcrQPSLimit, _ = strconv.Atoi(os.Getenv("TCR_QPS_LIMIT"))
-	if envTcrQPSLimit > 1 && envTcrQPSLimit < tcrQPSLimit {
-		tcrQPSLimit = envTcrQPSLimit
+	var envTcrQPSLimit, _ = strconv.Atoi(os.Getenv("REG_ADAPTER_TCR_QPS_LIMIT"))
+	if envTcrQPSLimit > tcrQPSLimit || envTcrQPSLimit < 1 {
+		envTcrQPSLimit = tcrQPSLimit
 	}
+	rateLimiterTransport = lib.NewRateLimitedTransport(envTcrQPSLimit, commonhttp.NewTransport())
 
 	if err := adp.RegisterFactory(model.RegistryTypeTencentTcr, new(factory)); err != nil {
 		log.Errorf("failed to register factory for %s: %v", model.RegistryTypeTencentTcr, err)
@@ -154,10 +157,7 @@ func newAdapter(registry *model.Registry) (a *adapter, err error) {
 	client.Init(*instanceInfo.RegionName).
 		WithCredential(tcrCredential).
 		WithProfile(cfp).
-		WithHttpTransport(newRateLimitedTransport(tcrQPSLimit, http.DefaultTransport))
-	if err != nil {
-		return
-	}
+		WithHttpTransport(rateLimiterTransport)
 
 	var credential = NewAuth(instanceInfo.RegistryId, client)
 	var transport = commonhttp.GetHTTPTransport(commonhttp.WithInsecure(registry.Insecure))

src/pkg/reg/adapter/tencentcr/artifact_registry.go

@@ -26,7 +26,7 @@ import (
 	"github.com/goharbor/harbor/src/pkg/reg/util"
 )
 
-var (
+const (
 	tcrQPSLimit = 10
 )