Line data    Source code

       1             : #pragma once
       2             : 
       3             : #include <cstdint>
       4             : #include <memory>
       5             : 
       6             : #include "envoy/common/pure.h"
       7             : #include "envoy/common/time.h"
       8             : 
       9             : namespace Envoy {
      10             : 
      11             : /**
      12             :  * This class defines an interface for the token bucket algorithm.
      13             :  *
      14             :  * https://en.wikipedia.org/wiki/Token_bucket
      15             :  */
      16             : class TokenBucket {
      17             : public:
      18           1 :   virtual ~TokenBucket() = default;
      19             : 
      20             :   /**
      21             :    * @param tokens supplies the number of tokens to be consumed.
      22             :    * @param allow_partial supplies whether the token bucket will allow consumption of less tokens
      23             :    *                      than asked for. If allow_partial is true, the bucket contains 3 tokens,
      24             :    *                      and the caller asks for 5, the bucket will return 3 tokens and now be
      25             :    *                      empty.
      26             :    * @return the number of tokens actually consumed.
      27             :    */
      28             :   virtual uint64_t consume(uint64_t tokens, bool allow_partial) PURE;
      29             : 
      30             :   /**
      31             :    * @param tokens supplies the number of tokens to be consumed.
      32             :    * @param allow_partial supplies whether the token bucket will allow consumption of less tokens
      33             :    *                      than asked for. If allow_partial is true, the bucket contains 3 tokens,
      34             :    *                      and the caller asks for 5, the bucket will return 3 tokens and now be
      35             :    *                      empty.
      36             :    * @param time_to_next_token out param indicating the approx time until next token is available.
      37             :    * @return the number of tokens actually consumed.
      38             :    */
      39             :   virtual uint64_t consume(uint64_t tokens, bool allow_partial,
      40             :                            std::chrono::milliseconds& time_to_next_token) PURE;
      41             : 
      42             :   /**
      43             :    * @return returns the approximate time until a next token is available. Currently it
      44             :    * returns the upper bound on the amount of time until a next token is available.
      45             :    */
      46             :   virtual std::chrono::milliseconds nextTokenAvailable() PURE;
      47             : 
      48             :   /**
      49             :    * Reset the bucket with a specific number of tokens. Refill will begin again from the time that
      50             :    * this routine is called.
      51             :    * Note: The reset call might be honored only the first time this method is called. Check the
      52             :    * concrete implementation to confirm.
      53             :    */
      54             :   virtual void maybeReset(uint64_t num_tokens) PURE;
      55             : };
      56             : 
      57             : using TokenBucketPtr = std::unique_ptr<TokenBucket>;
      58             : 
      59             : }; // namespace Envoy