add leaky bucket throttler

make generic cell rate throttler atomic

Author: 1pkg

README.MD

@@ -73,7 +73,7 @@ func WithPriority(ctx context.Context, priority uint8) context.Context
 func WithWeight(ctx context.Context, weight int64) context.Context
 // WithKey adds the provided key to the provided context
 // to add additional call identifier to context.
-// Resulted context is used by: `before`, `after`, `timed`, `adaptive`, `semaphore` and `cellrate` throtttlers.
+// Resulted context is used by: `before`, `after`, `timed`, `adaptive`, `semaphore`, `cellrate` and `bucket` throtttlers.
 func WithKey(ctx context.Context, key string) context.Context
 // WithMessage adds the provided message to the provided context
 // to add additional message that need to be used to context.
@@ -167,6 +167,7 @@ You can find list of returning error types for all existing throttlers in thrott
 | generator | `func NewThrottlerGenerator(gen Generator, capacity uint64, eviction float64) Throttler` | Creates new throttler instance that throttles if found key matching throttler throttles.<br> If no key matching throttler has been found generator used insted to provide new throttler that will be added to existing throttlers map.<br> Generated throttlers are kept in bounded map with capacity *c* defined by the specified capacity and eviction rate *e* defined by specified eviction value is normalized to [0.0, 1.0], where eviction rate affects number of throttlers that will be removed from the map after bounds overflow.<br> Use `WithKey` to specify key for throttler matching and generation.<br> - could return `ErrorInternal`;<br> - could return any underlying throttler error; |
 | semaphore | `func NewThrottlerSemaphore(weight int64) Throttler` | Creates new throttler instance that throttles call if underlying semaphore throttles.<br>Use `WithWeight` to override context call weight, 1 by default.<br> - could return `ErrorThreshold`; |
 | cellrate | `func NewThrottlerCellRate(threshold uint64, interval time.Duration, monotone bool) Throttler` | Creates new throttler instance that uses generic cell rate algorithm to throttles call within provided interval and threshold.<br>If provided monotone flag is set class to release will have no effect on throttler.<br>Use `WithWeight` to override context call qunatity, 1 by default.<br> - could return `ErrorThreshold`; |
+| bucket | `func NewThrottlerBucket(threshold uint64, interval time.Duration, monotone bool) Throttler` | Creates new throttler instance that leaky bucket algorithm to throttles call within provided interval and threshold.<br>If provided monotone flag is set class to release will have no effect on throttler.<br>Use `WithWeight` to override context call qunatity, 1 by default.<br> - could return `ErrorThreshold`; |
 
 ## Integrations
 

context.go

@@ -50,7 +50,7 @@ func ctxPriority(ctx context.Context, limit uint8) uint8 {
 
 // WithWeight adds the provided weight to the provided context
 // to differ `Acquire` weight levels.
-// Resulted context is used by: `before`, `after`, `timed`, `adaptive`, `semaphore` and `cellrate` throtttlers.
+// Resulted context is used by: `before`, `after`, `timed`, `adaptive`, `semaphore`, `cellrate` and `bucket` throtttlers.
 func WithWeight(ctx context.Context, weight int64) context.Context {
 	return context.WithValue(ctx, ghctxweight, weight)
 }
@@ -64,6 +64,15 @@ func ctxWeight(ctx context.Context) int64 {
 	return 1
 }
 
+func ctxWeightMod(ctx context.Context) int64 {
+	if val := ctx.Value(ghctxweight); val != nil {
+		if weight, ok := val.(int64); ok && weight > 0 {
+			return weight
+		}
+	}
+	return 1
+}
+
 // WithKey adds the provided key to the provided context
 // to add additional call identifier to context.
 // Resulted context is used by: `pattern` and `generator` throtttlers.

throttlers.go

@@ -1095,7 +1095,7 @@ func NewThrottlerSemaphore(weight int64) Throttler {
 }
 
 func (thr tsemaphore) Acquire(ctx context.Context) error {
-	if ok := thr.sem.TryAcquire(ctxWeight(ctx)); !ok {
+	if ok := thr.sem.TryAcquire(ctxWeightMod(ctx)); !ok {
 		return ErrorThreshold{
 			Throttler: "semaphore",
 			Threshold: strbool(ok),
@@ -1107,7 +1107,7 @@ func (thr tsemaphore) Acquire(ctx context.Context) error {
 func (thr tsemaphore) Release(ctx context.Context) error {
 	// prevent over releasing panic.
 	defer func() { _ = recover() }()
-	thr.sem.Release(ctxWeight(ctx))
+	thr.sem.Release(ctxWeightMod(ctx))
 	return nil
 }
 
@@ -1129,21 +1129,20 @@ func NewThrottlerCellRate(threshold uint64, interval time.Duration, monotone boo
 }
 
 func (thr *tcellrate) Acquire(ctx context.Context) error {
-	current := atomicGet(&thr.current)
 	nowTs := uint64(time.Now().UTC().UnixNano())
-	if current < nowTs {
-		current = nowTs
+	delta := (uint64(thr.quantum) * uint64(ctxWeightMod(ctx)))
+	if current := atomicGet(&thr.current); current < nowTs {
+		delta += nowTs - current
 	}
-	updated := current + (uint64(thr.quantum) * uint64(ctxWeight(ctx)))
 	max := nowTs + (uint64(thr.quantum) * thr.threshold)
-	if updated > max {
-		current := uint64(math.Round(float64(updated-nowTs) / float64(thr.quantum)))
+	if current := atomicBAdd(&thr.current, delta); current > max {
+		atomicBSub(&thr.current, delta)
+		cur := thr.threshold + uint64(math.Ceil(float64(current-max)/float64(thr.quantum)))
 		return ErrorThreshold{
 			Throttler: "cellrate",
-			Threshold: strpair{current: current, threshold: thr.threshold},
+			Threshold: strpair{current: cur, threshold: thr.threshold},
 		}
 	}
-	atomicSet(&thr.current, updated)
 	return nil
 }
 
@@ -1152,7 +1151,52 @@ func (thr *tcellrate) Release(ctx context.Context) error {
 	if thr.monotone {
 		return nil
 	}
-	updated := atomicGet(&thr.current) - (uint64(thr.quantum) * uint64(ctxWeight(ctx)))
-	atomicSet(&thr.current, updated)
+	atomicBSub(&thr.current, uint64(thr.quantum)*uint64(ctxWeightMod(ctx)))
+	return nil
+}
+
+type tbucket struct {
+	current   uint64
+	lastTs    uint64
+	threshold uint64
+	quantum   time.Duration
+	monotone  bool
+}
+
+// NewThrottlerBucket creates new throttler instance that
+// uses leaky bucket algorithm to throttles call within provided interval and threshold.
+// If provided monotone flag is set class to release will have no effect on throttler.
+// Use `WithWeight` to override context call qunatity, 1 by default.
+// - could return `ErrorThreshold`;
+func NewThrottlerBucket(threshold uint64, interval time.Duration, monotone bool) Throttler {
+	quantum := time.Duration(math.Ceil(float64(interval) / float64(threshold)))
+	return &tbucket{threshold: threshold, quantum: quantum, monotone: monotone}
+}
+
+func (thr *tbucket) Acquire(ctx context.Context) error {
+	nowTs := uint64(time.Now().UTC().UnixNano())
+	var delta int64
+	if lastTs := atomicGet(&thr.lastTs); lastTs > 0 {
+		delta = int64(float64(nowTs-lastTs) / float64(thr.quantum))
+		nowTs = uint64(delta) * uint64(thr.quantum)
+	}
+	diff := ctxWeightMod(ctx) - delta
+	if current := atomicBSingAdd(&thr.current, diff); current > thr.threshold {
+		atomicBSingAdd(&thr.current, -diff)
+		return ErrorThreshold{
+			Throttler: "bucket",
+			Threshold: strpair{current: current, threshold: thr.threshold},
+		}
+	}
+	atomicBAdd(&thr.lastTs, nowTs)
+	return nil
+}
+
+func (thr *tbucket) Release(ctx context.Context) error {
+	// don't decrement calls for monotone cell.
+	if thr.monotone {
+		return nil
+	}
+	atomicBSub(&thr.current, uint64(ctxWeightMod(ctx)))
 	return nil
 }

throttlers_test.go

@@ -20,11 +20,11 @@ const (
 	ms3_0  time.Duration = 3 * time.Millisecond
 	ms4_0  time.Duration = 4 * time.Millisecond
 	ms5_0  time.Duration = 5 * time.Millisecond
-	ms6_0  time.Duration = 6 * time.Millisecond
 	ms7_0  time.Duration = 7 * time.Millisecond
 	ms8_0  time.Duration = 8 * time.Millisecond
 	ms9_0  time.Duration = 9 * time.Millisecond
 	ms10_0 time.Duration = 10 * time.Millisecond
+	ms12_0 time.Duration = 12 * time.Millisecond
 	ms30_0 time.Duration = 30 * time.Millisecond
 )
 
@@ -1255,7 +1255,7 @@ func TestThrottlers(t *testing.T) {
 		},
 		"Throttler monotone cellrate should throttle on threshold": {
 			tms: 5,
-			thr: NewThrottlerCellRate(2, ms6_0, true),
+			thr: NewThrottlerCellRate(2, ms12_0, true),
 			ctxs: []context.Context{
 				context.TODO(),
 				context.TODO(),
@@ -1267,8 +1267,8 @@ func TestThrottlers(t *testing.T) {
 				nope,
 				nope,
 				nope,
-				delayed(ms4_0, nope),
-				delayed(ms3_0, nope),
+				delayed(ms7_0, nope),
+				delayed(ms9_0, nope),
 			},
 			errs: []error{
 				nil,
@@ -1280,7 +1280,7 @@ func TestThrottlers(t *testing.T) {
 				nil,
 				ErrorThreshold{
 					Throttler: "cellrate",
-					Threshold: strpair{current: 3, threshold: 2},
+					Threshold: strpair{current: 4, threshold: 2},
 				},
 			},
 		},
@@ -1315,6 +1315,68 @@ func TestThrottlers(t *testing.T) {
 				nil,
 			},
 		},
+		"Throttler monotone bucket should throttle on threshold": {
+			tms: 5,
+			thr: NewThrottlerBucket(2, ms12_0, true),
+			ctxs: []context.Context{
+				context.TODO(),
+				context.TODO(),
+				context.TODO(),
+				context.TODO(),
+				WithWeight(context.TODO(), 2),
+			},
+			pres: []Runnable{
+				nope,
+				nope,
+				nope,
+				delayed(ms7_0, nope),
+				delayed(ms9_0, nope),
+			},
+			errs: []error{
+				nil,
+				nil,
+				ErrorThreshold{
+					Throttler: "bucket",
+					Threshold: strpair{current: 3, threshold: 2},
+				},
+				nil,
+				ErrorThreshold{
+					Throttler: "bucket",
+					Threshold: strpair{current: 4, threshold: 2},
+				},
+			},
+		},
+		"Throttler not monotone bucket should throttle on threshold": {
+			tms: 5,
+			thr: NewThrottlerBucket(2, ms9_0, false),
+			acts: []Runnable{
+				delayed(ms5_0, nope),
+				delayed(ms5_0, nope),
+				delayed(ms5_0, nope),
+				nope,
+				nope,
+			},
+			pres: []Runnable{
+				nope,
+				nope,
+				nope,
+				nope,
+				delayed(ms7_0, nope),
+			},
+			errs: []error{
+				nil,
+				nil,
+				ErrorThreshold{
+					Throttler: "bucket",
+					Threshold: strpair{current: 3, threshold: 2},
+				},
+				ErrorThreshold{
+					Throttler: "bucket",
+					Threshold: strpair{current: 3, threshold: 2},
+				},
+				nil,
+			},
+		},
 	}
 	for tname, ptrtcase := range table {
 		t.Run(tname, func(t *testing.T) {