Line data    Source code

       1             : #pragma once
       2             : 
       3             : #include <chrono>
       4             : 
       5             : #include "envoy/common/time.h"
       6             : 
       7             : #include "source/common/common/lock_guard.h"
       8             : #include "source/common/common/thread.h"
       9             : 
      10             : namespace Envoy {
      11             : namespace Extensions {
      12             : namespace Tracers {
      13             : namespace XRay {
      14             : 
      15             : /**
      16             :  * Simple token-bucket algorithm that enables counting samples/traces used per second.
      17             :  */
      18             : class Reservoir {
      19             : public:
      20             :   /**
      21             :    * Creates a new reservoir that allows up to |traces_per_second| samples.
      22             :    */
      23           0 :   explicit Reservoir(uint32_t traces_per_second) : traces_per_second_(traces_per_second) {}
      24             : 
      25             :   Reservoir(const Reservoir& other)
      26             :       : traces_per_second_(other.traces_per_second_), used_(other.used_),
      27           0 :         time_point_(other.time_point_) {}
      28             : 
      29           0 :   Reservoir& operator=(const Reservoir& other) {
      30           0 :     if (this == &other) {
      31           0 :       return *this;
      32           0 :     }
      33           0 :     traces_per_second_ = other.traces_per_second_;
      34           0 :     used_ = other.used_;
      35           0 :     time_point_ = other.time_point_;
      36           0 :     return *this;
      37           0 :   }
      38             : 
      39             :   /**
      40             :    * Determines whether all samples have been used up for this particular second.
      41             :    * Every second, this reservoir starts over with a full bucket.
      42             :    *
      43             :    * @param now Used to compare against the last recorded time to determine if it's still within the
      44             :    * same second.
      45             :    */
      46           0 :   bool take(Envoy::MonotonicTime now) {
      47           0 :     Envoy::Thread::LockGuard lg(sync_);
      48           0 :     const auto diff = now - time_point_;
      49           0 :     if (diff > std::chrono::seconds(1)) {
      50           0 :       used_ = 0;
      51           0 :       time_point_ = now;
      52           0 :     }
      53             : 
      54           0 :     if (used_ >= traces_per_second_) {
      55           0 :       return false;
      56           0 :     }
      57             : 
      58           0 :     ++used_;
      59           0 :     return true;
      60           0 :   }
      61             : 
      62             : private:
      63             :   uint32_t traces_per_second_;
      64             :   uint32_t used_{0};
      65             :   Envoy::MonotonicTime time_point_;
      66             :   Envoy::Thread::MutexBasicLockable sync_;
      67             : };
      68             : 
      69             : } // namespace XRay
      70             : } // namespace Tracers
      71             : } // namespace Extensions
      72             : } // namespace Envoy