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
69
  explicit Reservoir(uint32_t traces_per_second) : traces_per_second_(traces_per_second) {}
24

            
25
  Reservoir(const Reservoir& other)
26
59
      : traces_per_second_(other.traces_per_second_), used_(other.used_),
27
59
        time_point_(other.time_point_) {}
28

            
29
25
  Reservoir& operator=(const Reservoir& other) {
30
25
    if (this == &other) {
31
      return *this;
32
    }
33
25
    traces_per_second_ = other.traces_per_second_;
34
25
    used_ = other.used_;
35
25
    time_point_ = other.time_point_;
36
25
    return *this;
37
25
  }
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
77
  bool take(Envoy::MonotonicTime now) {
47
77
    Envoy::Thread::LockGuard lg(sync_);
48
77
    const auto diff = now - time_point_;
49
77
    if (diff > std::chrono::seconds(1)) {
50
5
      used_ = 0;
51
5
      time_point_ = now;
52
5
    }
53

            
54
77
    if (used_ >= traces_per_second_) {
55
70
      return false;
56
70
    }
57

            
58
7
    ++used_;
59
7
    return true;
60
77
  }
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