|
| 1 | +import 'dart:async'; |
| 2 | +import 'dart:math'; |
| 3 | + |
| 4 | +import 'package:clock/clock.dart'; |
| 5 | +import 'package:token_bucket_algorithm/src/token_bucket_storage.dart'; |
| 6 | + |
| 7 | +part 'token_bucket_state.dart'; |
| 8 | + |
| 9 | +abstract class _BaseTokenBucket<S extends AsyncTokenBucketStorage> { |
| 10 | + /// The maximum token amount of this bucket. |
| 11 | + final int size; |
| 12 | + |
| 13 | + /// The interval for refilling the bucket. |
| 14 | + final Duration refillInterval; |
| 15 | + |
| 16 | + /// The amount of tokens that is refilled every [refillInterval]. |
| 17 | + final int refillAmount; |
| 18 | + |
| 19 | + /// The storage for the internal [TokenBucketState]. |
| 20 | + final S storage; |
| 21 | + |
| 22 | + _BaseTokenBucket({ |
| 23 | + required this.size, |
| 24 | + required this.refillInterval, |
| 25 | + required this.refillAmount, |
| 26 | + required this.storage, |
| 27 | + }) : assert(size > 0), |
| 28 | + assert(refillAmount > 0), |
| 29 | + assert(refillInterval > Duration.zero); |
| 30 | + |
| 31 | + /// Returns the currently available tokens of this bucket. |
| 32 | + FutureOr<int> get availableTokens; |
| 33 | + |
| 34 | + /// Consumes [cost] tokens and returns whether the consuming was successful. |
| 35 | + FutureOr<bool> consume([int cost = 1]); |
| 36 | + |
| 37 | + TokenBucketState _refillBucket(TokenBucketState state) { |
| 38 | + final now = clock.now(); |
| 39 | + |
| 40 | + // relevant if e.g. user updates time in settings |
| 41 | + if (state.lastRefillTime.isAfter(now)) { |
| 42 | + return state.copyWith(lastRefillTime: now); |
| 43 | + } |
| 44 | + |
| 45 | + final refillTimes = (now.difference(state.lastRefillTime)).inMicroseconds ~/ |
| 46 | + refillInterval.inMicroseconds; |
| 47 | + |
| 48 | + final newTokenCount = min(size, state.tokens + refillTimes * refillAmount); |
| 49 | + |
| 50 | + return state.copyWith( |
| 51 | + tokens: newTokenCount, |
| 52 | + lastRefillTime: state.lastRefillTime.add(refillInterval * refillTimes), |
| 53 | + ); |
| 54 | + } |
| 55 | +} |
| 56 | + |
| 57 | +/// An async token bucket which can have an async [storage]. |
| 58 | +class AsyncTokenBucket extends _BaseTokenBucket<AsyncTokenBucketStorage> { |
| 59 | + Future<void> _future = Future.value(); |
| 60 | + |
| 61 | + AsyncTokenBucket({ |
| 62 | + required super.size, |
| 63 | + required super.refillInterval, |
| 64 | + required super.refillAmount, |
| 65 | + AsyncTokenBucketStorage? storage, |
| 66 | + }) : super(storage: storage ?? MemoryTokenBucketStorage()); |
| 67 | + |
| 68 | + @override |
| 69 | + FutureOr<int> get availableTokens async { |
| 70 | + await _queueFuture(() async { |
| 71 | + await storage.set(_refillBucket(await storage.get())); |
| 72 | + return false; |
| 73 | + }); |
| 74 | + return Future.value(storage.get()).then((state) => state.tokens); |
| 75 | + } |
| 76 | + |
| 77 | + @override |
| 78 | + FutureOr<bool> consume([int cost = 1]) async { |
| 79 | + if (cost < 1 || cost > size) { |
| 80 | + throw ArgumentError('cost must be <=$size and >=1'); |
| 81 | + } |
| 82 | + return _queueFuture(() async { |
| 83 | + final state = _refillBucket(await storage.get()); |
| 84 | + final (result, newState) = state.consume(cost); |
| 85 | + await storage.set(newState); |
| 86 | + return result; |
| 87 | + }); |
| 88 | + } |
| 89 | + |
| 90 | + Future<T> _queueFuture<T>(Future<T> Function() computation) { |
| 91 | + final newFuture = _future.then((_) => computation()); |
| 92 | + _future = newFuture.whenComplete(() { |
| 93 | + if (_future == newFuture) _future = Future.value(); |
| 94 | + }); |
| 95 | + return newFuture; |
| 96 | + } |
| 97 | +} |
| 98 | + |
| 99 | +/// A standard token bucket. |
| 100 | +/// |
| 101 | +/// If you want to store the [TokenBucketState] in an [AsyncTokenBucketStorage], |
| 102 | +/// you have to use [AsyncTokenBucket] instead. |
| 103 | +class TokenBucket extends _BaseTokenBucket<TokenBucketStorage> { |
| 104 | + TokenBucket({ |
| 105 | + required super.size, |
| 106 | + required super.refillInterval, |
| 107 | + required super.refillAmount, |
| 108 | + TokenBucketStorage? storage, |
| 109 | + }) : super(storage: storage ?? MemoryTokenBucketStorage()); |
| 110 | + |
| 111 | + @override |
| 112 | + int get availableTokens { |
| 113 | + storage.set(_refillBucket(storage.get())); |
| 114 | + return storage.get().tokens; |
| 115 | + } |
| 116 | + |
| 117 | + @override |
| 118 | + bool consume([int cost = 1]) { |
| 119 | + if (cost < 1 || cost > size) { |
| 120 | + throw ArgumentError('cost must be <=$size and >=1'); |
| 121 | + } |
| 122 | + final state = _refillBucket(storage.get()); |
| 123 | + final (result, newState) = state.consume(cost); |
| 124 | + storage.set(newState); |
| 125 | + return result; |
| 126 | + } |
| 127 | +} |
0 commit comments