Skip to content

Commit e57f97f

Browse files
Fix token count in archival rate limiter (#3749)
1 parent 92b8758 commit e57f97f

File tree

2 files changed

+4
-3
lines changed

2 files changed

+4
-3
lines changed

service/history/archival/archiver.go

+3-2
Original file line numberDiff line numberDiff line change
@@ -143,14 +143,15 @@ func (a *archiver) Archive(ctx context.Context, request *Request) (res *Response
143143
metricsScope.Timer(metrics.ArchiverArchiveLatency.GetMetricName()).
144144
Record(time.Since(start), metrics.StringTag("status", status))
145145
}(time.Now())
146-
if err := a.rateLimiter.WaitN(ctx, 2); err != nil {
146+
numTargets := len(request.Targets)
147+
if err := a.rateLimiter.WaitN(ctx, numTargets); err != nil {
147148
return nil, &serviceerror.ResourceExhausted{
148149
Cause: enumspb.RESOURCE_EXHAUSTED_CAUSE_RPS_LIMIT,
149150
Message: fmt.Sprintf("archival rate limited: %s", err.Error()),
150151
}
151152
}
152153
var wg sync.WaitGroup
153-
errs := make([]error, len(request.Targets))
154+
errs := make([]error, numTargets)
154155
for i, target := range request.Targets {
155156
wg.Add(1)
156157
i := i

service/history/archival/archiver_test.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -280,7 +280,7 @@ func TestArchiver(t *testing.T) {
280280
Return(c.ArchiveVisibilityErr)
281281
}
282282
rateLimiter := quotas.NewMockRateLimiter(controller)
283-
rateLimiter.EXPECT().WaitN(gomock.Any(), 2).Return(c.RateLimiterWaitErr)
283+
rateLimiter.EXPECT().WaitN(gomock.Any(), len(c.Targets)).Return(c.RateLimiterWaitErr)
284284

285285
// we need this channel to get the Archiver which is created asynchronously
286286
archivers := make(chan Archiver, 1)

0 commit comments

Comments
 (0)