ratelimit: set cancellation earlier (#402)

future may be completed by worker before cancellation is set
This commit is contained in:
Jacek Sieka 2023-06-05 13:47:38 +02:00 committed by GitHub
parent a6ac5f2213
commit 2ef34c7339
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 6 additions and 4 deletions

View File

@ -102,10 +102,6 @@ proc consume*(bucket: TokenBucket, tokens: int): Future[void] =
retFuture.complete() retFuture.complete()
return retFuture return retFuture
bucket.pendingRequests.add(BucketWaiter(future: retFuture, value: tokens))
if isNil(bucket.workFuture) or bucket.workFuture.finished():
bucket.workFuture = worker(bucket)
proc cancellation(udata: pointer) = proc cancellation(udata: pointer) =
for index in 0..<bucket.pendingRequests.len: for index in 0..<bucket.pendingRequests.len:
if bucket.pendingRequests[index].future == retFuture: if bucket.pendingRequests[index].future == retFuture:
@ -115,6 +111,12 @@ proc consume*(bucket: TokenBucket, tokens: int): Future[void] =
bucket.manuallyReplenished.fire() bucket.manuallyReplenished.fire()
break break
retFuture.cancelCallback = cancellation retFuture.cancelCallback = cancellation
bucket.pendingRequests.add(BucketWaiter(future: retFuture, value: tokens))
if isNil(bucket.workFuture) or bucket.workFuture.finished():
bucket.workFuture = worker(bucket)
return retFuture return retFuture
proc replenish*(bucket: TokenBucket, tokens: int) = proc replenish*(bucket: TokenBucket, tokens: int) =