diff --git a/Source/Esendex.TokenBucket.Tests/TokenBucketTests.cs b/Source/Esendex.TokenBucket.Tests/TokenBucketTests.cs index 59d968a..045c1de 100644 --- a/Source/Esendex.TokenBucket.Tests/TokenBucketTests.cs +++ b/Source/Esendex.TokenBucket.Tests/TokenBucketTests.cs @@ -87,6 +87,15 @@ public void ConsumeWhenTokenAvailable() _sleepStrategy.Verify(s => s.Sleep(), Times.Never()); } + [Test, Timeout(ConsumeTimeout)] + public void ConsumeAsyncWhenTokenAvailable() + { + _refillStrategy.AddToken(); + _bucket.ConsumeAsync().Wait(); + + _sleepStrategy.Verify(s => s.Sleep(), Times.Never()); + } + [Test, Timeout(ConsumeTimeout)] public void ConsumeWhenTokensAvailable() { @@ -97,6 +106,16 @@ public void ConsumeWhenTokensAvailable() _sleepStrategy.Verify(s => s.Sleep(), Times.Never()); } + [Test, Timeout(ConsumeTimeout)] + public void ConsumeAsyncWhenTokensAvailable() + { + const int tokensToConsume = 2; + _refillStrategy.AddTokens(tokensToConsume); + _bucket.ConsumeAsync(tokensToConsume).Wait(); + + _sleepStrategy.Verify(s => s.Sleep(), Times.Never()); + } + [Test, Timeout(ConsumeTimeout)] public void ConsumeWhenTokenUnavailable() { @@ -110,6 +129,19 @@ public void ConsumeWhenTokenUnavailable() _sleepStrategy.Verify(); } + [Test, Timeout(ConsumeTimeout)] + public void ConsumeAsyncWhenTokenUnavailable() + { + _sleepStrategy + .Setup(s => s.Sleep()) + .Callback(_refillStrategy.AddToken) + .Verifiable(); + + _bucket.ConsumeAsync().Wait(); + + _sleepStrategy.Verify(); + } + [Test, Timeout(ConsumeTimeout)] public void ConsumeWhenTokensUnavailable() { @@ -124,6 +156,20 @@ public void ConsumeWhenTokensUnavailable() _sleepStrategy.Verify(); } + [Test, Timeout(ConsumeTimeout)] + public void ConsumeAsyncWhenTokensUnavailable() + { + const int tokensToConsume = 7; + _sleepStrategy + .Setup(s => s.Sleep()) + .Callback(() => _refillStrategy.AddTokens(tokensToConsume)) + .Verifiable(); + + _bucket.ConsumeAsync(tokensToConsume).Wait(); + + _sleepStrategy.Verify(); + } + private sealed class MockRefillStrategy : IRefillStrategy { private long _numTokensToAdd; diff --git a/Source/Esendex.TokenBucket/ITokenBucket.cs b/Source/Esendex.TokenBucket/ITokenBucket.cs index 63b2336..5a9e08c 100644 --- a/Source/Esendex.TokenBucket/ITokenBucket.cs +++ b/Source/Esendex.TokenBucket/ITokenBucket.cs @@ -1,3 +1,5 @@ +using System.Threading.Tasks; + namespace Esendex.TokenBucket { /// @@ -29,10 +31,23 @@ public interface ITokenBucket /// void Consume(); + /// + /// Consume a single token from the bucket asynchronously. This method does not block + /// A task that returns once a token has been consumed + /// + Task ConsumeAsync(); + /// /// Consumes multiple tokens from the bucket. If enough tokens are not currently available then this method will block /// /// The number of tokens to consume from the bucket, must be a positive number. void Consume(long numTokens); + + /// + /// Consume multiple tokens from the bucket asynchronously. This method does not block + /// The number of tokens to consume from the bucket, must be a positive number. + /// A task that returns once the requested tokens have been consumed + /// + Task ConsumeAsync(long numTokens); } } \ No newline at end of file diff --git a/Source/Esendex.TokenBucket/TokenBucket.cs b/Source/Esendex.TokenBucket/TokenBucket.cs index 0b22bcc..17a1ffa 100644 --- a/Source/Esendex.TokenBucket/TokenBucket.cs +++ b/Source/Esendex.TokenBucket/TokenBucket.cs @@ -1,4 +1,5 @@ using System; +using System.Threading.Tasks; namespace Esendex.TokenBucket { @@ -72,25 +73,44 @@ public bool TryConsume(long numTokens) /// /// Consume a single token from the bucket. If no token is currently available then this method will block until a /// token becomes available. - /// - public void Consume() - { - Consume(1); + /// + public void Consume() + { + Consume(1); + } + + /// + /// Consume a single token from the bucket asynchronously. This method does not block + /// A task that returns once a token has been consumed + /// + public Task ConsumeAsync() + { + return ConsumeAsync(1); } /// /// Consumes multiple tokens from the bucket. If enough tokens are not currently available then this method will block /// - /// The number of tokens to consume from the bucket, must be a positive number. - public void Consume(long numTokens) - { - while (true) { - if (TryConsume(numTokens)) { - break; - } - - _sleepStrategy.Sleep(); - } + /// The number of tokens to consume from the bucket, must be a positive number. + public void Consume(long numTokens) + { + while (true) { + if (TryConsume(numTokens)) { + break; + } + + _sleepStrategy.Sleep(); + } + } + + /// + /// Consume multiple tokens from the bucket asynchronously. This method does not block + /// The number of tokens to consume from the bucket, must be a positive number. + /// A task that returns once the requested tokens have been consumed + /// + public Task ConsumeAsync(long numTokens) + { + return Task.Factory.StartNew(() => Consume(numTokens)); } } }