From e57f97f9f831e3cd29ea014b4b6d606ddb24e004 Mon Sep 17 00:00:00 2001 From: Michael Snowden Date: Wed, 28 Dec 2022 10:19:32 -0800 Subject: [PATCH] Fix token count in archival rate limiter (#3749) --- service/history/archival/archiver.go | 5 +++-- service/history/archival/archiver_test.go | 2 +- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/service/history/archival/archiver.go b/service/history/archival/archiver.go index 3988b8cc626..07a079c537d 100644 --- 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 index 296b250a276..797480c5f23 100644 --- 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)