Fix token count in archival rate limiter (#3749)

temporalio · Dec 28, 2022 · e57f97f · e57f97f
1 parent 92b8758
commit e57f97f
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/service/history/archival/archiver.go b/service/history/archival/archiver.go
@@ -143,14 +143,15 @@ func (a *archiver) Archive(ctx context.Context, request *Request) (res *Response
 		metricsScope.Timer(metrics.ArchiverArchiveLatency.GetMetricName()).
 			Record(time.Since(start), metrics.StringTag("status", status))
 	}(time.Now())
-	if err := a.rateLimiter.WaitN(ctx, 2); err != nil {
+	numTargets := len(request.Targets)
+	if err := a.rateLimiter.WaitN(ctx, numTargets); err != nil {
 		return nil, &serviceerror.ResourceExhausted{
 			Cause:   enumspb.RESOURCE_EXHAUSTED_CAUSE_RPS_LIMIT,
 			Message: fmt.Sprintf("archival rate limited: %s", err.Error()),
 		}
 	}
 	var wg sync.WaitGroup
-	errs := make([]error, len(request.Targets))
+	errs := make([]error, numTargets)
 	for i, target := range request.Targets {
 		wg.Add(1)
 		i := i

diff --git a/service/history/archival/archiver_test.go b/service/history/archival/archiver_test.go
@@ -280,7 +280,7 @@ func TestArchiver(t *testing.T) {
 					Return(c.ArchiveVisibilityErr)
 			}
 			rateLimiter := quotas.NewMockRateLimiter(controller)
-			rateLimiter.EXPECT().WaitN(gomock.Any(), 2).Return(c.RateLimiterWaitErr)
+			rateLimiter.EXPECT().WaitN(gomock.Any(), len(c.Targets)).Return(c.RateLimiterWaitErr)
 
 			// we need this channel to get the Archiver which is created asynchronously
 			archivers := make(chan Archiver, 1)