From 601c7370bfd28833518efea45988dec8677227b9 Mon Sep 17 00:00:00 2001 From: Sami Kerola Date: Mon, 12 Sep 2022 16:05:44 +0100 Subject: [PATCH] add thread-safe tockenbucket touch() function Making the touch() thread-safe allows removing external mutex(es) that protect calling the touch. Unarguably the change also makes the data racing more, but that is fine. With low frequency updates the tocket bucket will not trip over, and with high frequency updates one should not be able to observe off-by-one type of return value distortion when values are stored few jiffies apart. Signed-off-by: Sami Kerola --- tokenbucket/bucket.go | 27 +++++++++++++++++---------- 1 file changed, 17 insertions(+), 10 deletions(-) diff --git a/tokenbucket/bucket.go b/tokenbucket/bucket.go index c901d93..80221a4 100644 --- a/tokenbucket/bucket.go +++ b/tokenbucket/bucket.go @@ -3,6 +3,7 @@ package tokenbucket import ( "math/rand" + "sync/atomic" "time" ) @@ -43,19 +44,25 @@ func New(num int, rate float64, depth uint64) *Filter { func (b *Filter) touch(it *item) bool { now := uint64(time.Now().UnixNano()) - delta := now - it.prev - it.credit += delta - it.prev = now + oldPrev := atomic.LoadUint64(&it.prev) + oldCredit := atomic.LoadUint64(&it.credit) - if it.credit > b.creditMax { - it.credit = b.creditMax - } + delta := now - oldPrev + defer atomic.StoreUint64(&it.prev, now) + + newCredit := oldCredit + delta + allow := false - if it.credit > b.touchCost { - it.credit -= b.touchCost - return true + if newCredit > b.creditMax { + newCredit = b.creditMax } - return false + if newCredit > b.touchCost { + newCredit -= b.touchCost + allow = true + } + atomic.StoreUint64(&it.credit, newCredit) + + return allow } // Touch finds the token bucket for d, takes a token out of it and reports if