throttle: add non-blocking get_or_fail method

1 files changed, 12 insertions, 1 deletions

diff --git a/src/common/Throttle.h b/src/common/Throttle.h
index e940937b9ea..cd772dd8b29 100644
--- a/src/common/Throttle.h
+++ b/src/common/Throttle.h
@@ -69,7 +69,7 @@ public:
   }
 
   bool get(int64_t c = 1, int64_t m = 0) {
-    assert(c > 0);
+    assert(c >= 0);
     Mutex::Locker l(lock);
     if (m) {
       assert(m > 0);
@@ -80,6 +80,17 @@ public:
     return waited;
   }
 
+  /* Returns true if it successfully got the requested amount,
+   * or false if it would block.
+   */
+  bool get_or_fail(int64_t c = 1) {
+    assert (c >= 0);
+    Mutex::Locker l(lock);
+    if (_should_wait(c)) return false;
+    count += c;
+    return true;
+  }
+
   int64_t put(int64_t c = 1) {
     assert(c >= 0);
     Mutex::Locker l(lock);