From 4393bdeb77399a81ac243af8ae510d244be52131 Mon Sep 17 00:00:00 2001
From: Klaus Basan <klaus.basan@klausbasan.de>
Date: Thu, 7 Jun 2018 01:16:24 +0200
Subject: [PATCH] Ref T275, CTokenBucket further improvements

* allow to use extrenal timestamp for replenishment
* only replenish if there are not enough tokens "in stock"
---
 src/blackmisc/tokenbucket.cpp | 33 +++++++++++++++++++++++++++------
 src/blackmisc/tokenbucket.h   | 25 ++++++++++++++++++-------
 2 files changed, 45 insertions(+), 13 deletions(-)

diff --git a/src/blackmisc/tokenbucket.cpp b/src/blackmisc/tokenbucket.cpp
index 53c91881f..0d78bddce 100644
--- a/src/blackmisc/tokenbucket.cpp
+++ b/src/blackmisc/tokenbucket.cpp
@@ -24,16 +24,25 @@ namespace BlackMisc
         : m_capacity(capacity), m_numTokensToRefill(numTokensToRefill), m_intervalMs(intervalMs)
     {}
 
-    bool CTokenBucket::tryConsume(int numTokens)
+    bool CTokenBucket::tryConsume(int numTokens, qint64 msSinceEpoch)
     {
         Q_ASSERT(numTokens > 0 && numTokens < m_capacity);
 
+        // enough tokens in stock?
+        if (numTokens <= m_availableTokens)
+        {
+            m_availableTokens -= numTokens;
+            return true;
+        }
+
         // Replenish maximal up to capacity
-        int replenishedTokens = qMin(m_capacity, this->getTokens());
+        const int tokens = this->getTokens(msSinceEpoch);
+        const int replenishedTokens = qMin(m_capacity, tokens);
 
         // Take care of overflows
         m_availableTokens = qMin(m_availableTokens + replenishedTokens, m_capacity);
 
+        // check again after replenishment
         if (numTokens <= m_availableTokens)
         {
             m_availableTokens -= numTokens;
@@ -42,9 +51,9 @@ namespace BlackMisc
         return false;
     }
 
-    void CTokenBucket::setNumberOfTokensToRefill(int noTokens)
+    void CTokenBucket::setNumberOfTokensToRefill(int numTokens)
     {
-        m_numTokensToRefill = noTokens;
+        m_numTokensToRefill = numTokens;
     }
 
     void CTokenBucket::setCapacity(int capacity)
@@ -52,9 +61,21 @@ namespace BlackMisc
         m_capacity = capacity;
     }
 
-    int CTokenBucket::getTokens()
+    void CTokenBucket::setCapacityAndTokensToRefill(int numTokens)
     {
-        const qint64 now = QDateTime::currentMSecsSinceEpoch();
+        this->setCapacity(numTokens);
+        this->setNumberOfTokensToRefill(numTokens);
+    }
+
+    int CTokenBucket::getTokensPerSecond() const
+    {
+        if (m_intervalMs < 1) { return 0; }
+        return m_numTokensToRefill * 1000 / m_intervalMs;
+    }
+
+    int CTokenBucket::getTokens(qint64 msSinceEpoch)
+    {
+        const qint64 now = msSinceEpoch > 0 ? msSinceEpoch : QDateTime::currentMSecsSinceEpoch();
         const qint64 deltaMs = now - m_lastReplenishmentTime;
         const int numberOfTokens = static_cast<int>(m_numTokensToRefill * deltaMs / m_intervalMs);
 
diff --git a/src/blackmisc/tokenbucket.h b/src/blackmisc/tokenbucket.h
index ef672e997..53d5ea5b5 100644
--- a/src/blackmisc/tokenbucket.h
+++ b/src/blackmisc/tokenbucket.h
@@ -33,25 +33,36 @@ namespace BlackMisc
         CTokenBucket(int capacity, qint64 intervalMs, int numTokensToRefill);
 
         //! Try to consume a number of tokens
-        bool tryConsume(int numTokens = 1);
+        //! \remark if a current timestamp is already available, it can be passed
+        bool tryConsume(int numTokens = 1, qint64 msSinceEpoch = -1);
 
         //! Number of tokens to refill
-        void setNumberOfTokensToRefill(int noTokens);
+        void setNumberOfTokensToRefill(int numTokens);
 
         //! Set the capacity
         void setCapacity(int capacity);
 
+        //! Tokens/capacity if both are same
+        void setCapacityAndTokensToRefill(int numTokens);
+
+        //! Set the interval
+        void setInterval(qint64 intervalMs) { m_intervalMs = intervalMs; }
+
+        //! Tokens per second
+        int getTokensPerSecond() const;
+
     private:
         //! Get available tokens since last replenishment.
         //! \note Note that replenishment is implemented lazy.
         //!    This means, tokens will not replenished on regular basis via a running timer,
         //!    but they will be replenished while trying to consume them.
-        int getTokens();
+        //! \remark a timestamp can be passed, or "now" is taken
+        int getTokens(qint64 msSinceEpoch);
 
-        int m_capacity = 10;        //!< Maximum capacity of tokens
-        int m_availableTokens = 10; //!< Currently available tokens. The initial value is 10
-        int m_numTokensToRefill;    //!< Number of tokens to be refilled each interval
-        qint64 m_intervalMs = 5000; //!< Refill interval, e.g. every 5 secs
+        int m_capacity = 10;         //!< Maximum capacity of tokens
+        int m_availableTokens = 10;  //!< Currently available tokens. The initial value is 10
+        int m_numTokensToRefill = 1; //!< Number of tokens to be refilled each interval
+        qint64 m_intervalMs = 5000;  //!< Refill interval, e.g. every 5 secs
         qint64 m_lastReplenishmentTime = QDateTime::currentMSecsSinceEpoch(); //!< Last time
     };
 } // ns