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));
}
}
}