src/lib/evloop/token_bucket.h

   1 /* Copyright (c) 2018-2021, The Tor Project, Inc. */
   2 /* See LICENSE for licensing information */
   3
   4 /**
   5  * \file token_bucket.h
   6  * \brief Headers for token_bucket.c
   7  **/
   8
   9 #ifndef TOR_TOKEN_BUCKET_H
  10 #define TOR_TOKEN_BUCKET_H
  11
  12 #include "lib/cc/torint.h"
  13 #include "lib/testsupport/testsupport.h"
  14
  15 /** Largest allowable burst value for a token buffer. */
  16 #define TOKEN_BUCKET_MAX_BURST INT32_MAX
  17
  18 /** A generic token buffer configuration: determines the number of tokens
  19  * added to the bucket in each time unit (the "rate"), and the maximum number
  20  * of tokens in the bucket (the "burst") */
  21 typedef struct token_bucket_cfg_t {
  22   uint32_t rate;
  23   int32_t burst;
  24 } token_bucket_cfg_t;
  25
  26 /** A raw token bucket, decoupled from its configuration and timestamp. */
  27 typedef struct token_bucket_raw_t {
  28   int32_t bucket;
  29 } token_bucket_raw_t;
  30
  31 void token_bucket_cfg_init(token_bucket_cfg_t *cfg,
  32                            uint32_t rate,
  33                            uint32_t burst);
  34
  35 void token_bucket_raw_adjust(token_bucket_raw_t *bucket,
  36                              const token_bucket_cfg_t *cfg);
  37
  38 void token_bucket_raw_reset(token_bucket_raw_t *bucket,
  39                             const token_bucket_cfg_t *cfg);
  40
  41 int token_bucket_raw_dec(token_bucket_raw_t *bucket,
  42                          ssize_t n);
  43
  44 int token_bucket_raw_refill_steps(token_bucket_raw_t *bucket,
  45                                   const token_bucket_cfg_t *cfg,
  46                                   const uint32_t elapsed_steps);
  47
  48 static inline size_t token_bucket_raw_get(const token_bucket_raw_t *bucket);
  49 /** Return the current number of bytes set in a token bucket. */
  50 static inline size_t
  51 token_bucket_raw_get(const token_bucket_raw_t *bucket)
  52 {
  53   return bucket->bucket >= 0 ? bucket->bucket : 0;
  54 }
  55
  56 /** A convenience type containing all the pieces needed for a coupled
  57  * read-bucket and write-bucket that have the same rate limit, and which use
  58  * "timestamp units" (see compat_time.h) for their time. */
  59 typedef struct token_bucket_rw_t {
  60   token_bucket_cfg_t cfg;
  61   token_bucket_raw_t read_bucket;
  62   token_bucket_raw_t write_bucket;
  63   uint32_t last_refilled_at_timestamp;
  64 } token_bucket_rw_t;
  65
  66 void token_bucket_rw_init(token_bucket_rw_t *bucket,
  67                           uint32_t rate,
  68                           uint32_t burst,
  69                           uint32_t now_ts);
  70
  71 void token_bucket_rw_adjust(token_bucket_rw_t *bucket,
  72                             uint32_t rate, uint32_t burst);
  73
  74 void token_bucket_rw_reset(token_bucket_rw_t *bucket,
  75                            uint32_t now_ts);
  76
  77 #define TB_READ 1
  78 #define TB_WRITE 2
  79
  80 int token_bucket_rw_refill(token_bucket_rw_t *bucket,
  81                            uint32_t now_ts);
  82
  83 int token_bucket_rw_dec_read(token_bucket_rw_t *bucket,
  84                              ssize_t n);
  85 int token_bucket_rw_dec_write(token_bucket_rw_t *bucket,
  86                               ssize_t n);
  87
  88 int token_bucket_rw_dec(token_bucket_rw_t *bucket,
  89                         ssize_t n_read, ssize_t n_written);
  90
  91 static inline size_t token_bucket_rw_get_read(const token_bucket_rw_t *bucket);
  92 static inline size_t
  93 token_bucket_rw_get_read(const token_bucket_rw_t *bucket)
  94 {
  95   return token_bucket_raw_get(&bucket->read_bucket);
  96 }
  97
  98 static inline size_t token_bucket_rw_get_write(
  99                                             const token_bucket_rw_t *bucket);
 100 static inline size_t
 101 token_bucket_rw_get_write(const token_bucket_rw_t *bucket)
 102 {
 103   return token_bucket_raw_get(&bucket->write_bucket);
 104 }
 105
 106 /**
 107  * A specialized bucket containing a single counter.
 108  */
 109
 110 typedef struct token_bucket_ctr_t {
 111   token_bucket_cfg_t cfg;
 112   token_bucket_raw_t counter;
 113   uint32_t last_refilled_at_timestamp;
 114 } token_bucket_ctr_t;
 115
 116 void token_bucket_ctr_init(token_bucket_ctr_t *bucket, uint32_t rate,
 117                            uint32_t burst, uint32_t now_ts);
 118 void token_bucket_ctr_adjust(token_bucket_ctr_t *bucket, uint32_t rate,
 119                              uint32_t burst);
 120 void token_bucket_ctr_reset(token_bucket_ctr_t *bucket, uint32_t now_ts);
 121 void token_bucket_ctr_refill(token_bucket_ctr_t *bucket, uint32_t now_ts);
 122
 123 static inline bool
 124 token_bucket_ctr_dec(token_bucket_ctr_t *bucket, ssize_t n)
 125 {
 126   return token_bucket_raw_dec(&bucket->counter, n);
 127 }
 128
 129 static inline size_t
 130 token_bucket_ctr_get(const token_bucket_ctr_t *bucket)
 131 {
 132   return token_bucket_raw_get(&bucket->counter);
 133 }
 134
 135 #ifdef TOKEN_BUCKET_PRIVATE
 136
 137 /* To avoid making the rates too small, we consider units of "steps",
 138  * where a "step" is defined as this many timestamp ticks.  Keep this
 139  * a power of two if you can. */
 140 #define TICKS_PER_STEP 16
 141
 142 STATIC uint32_t rate_per_sec_to_rate_per_step(uint32_t rate);
 143
 144 #endif /* defined(TOKEN_BUCKET_PRIVATE) */
 145
 146 #endif /* !defined(TOR_TOKEN_BUCKET_H) */