Tor  0.4.7.0-alpha-dev
token_bucket.c
Go to the documentation of this file.
1 /* Copyright (c) 2018-2021, The Tor Project, Inc. */
2 /* See LICENSE for licensing information */
3 
4 /**
5  * \file token_bucket.c
6  * \brief Functions to use and manipulate token buckets, used for
7  * rate-limiting on connections and globally.
8  *
9  * Tor uses these token buckets to keep track of bandwidth usage, and
10  * sometimes other things too.
11  *
12  * There are two layers of abstraction here: "raw" token buckets, in which all
13  * the pieces are decoupled, and "read-write" token buckets, which combine all
14  * the moving parts into one.
15  *
16  * Token buckets may become negative.
17  **/
18 
19 #define TOKEN_BUCKET_PRIVATE
20 
22 #include "lib/log/util_bug.h"
23 #include "lib/intmath/cmp.h"
24 #include "lib/time/compat_time.h"
25 
26 #include <string.h>
27 
28 /**
29  * Set the <b>rate</b> and <b>burst</b> value in a token_bucket_cfg.
30  *
31  * Note that the <b>rate</b> value is in arbitrary units, but those units will
32  * determine the units of token_bucket_raw_dec(), token_bucket_raw_refill, and
33  * so on.
34  */
35 void
37  uint32_t rate,
38  uint32_t burst)
39 {
40  tor_assert_nonfatal(rate > 0);
41  tor_assert_nonfatal(burst > 0);
42  if (burst > TOKEN_BUCKET_MAX_BURST)
43  burst = TOKEN_BUCKET_MAX_BURST;
44 
45  cfg->rate = rate;
46  cfg->burst = burst;
47 }
48 
49 /**
50  * Initialize a raw token bucket and its associated timestamp to the "full"
51  * state, according to <b>cfg</b>.
52  */
53 void
55  const token_bucket_cfg_t *cfg)
56 {
57  bucket->bucket = cfg->burst;
58 }
59 
60 /**
61  * Adust a preexisting token bucket to respect the new configuration
62  * <b>cfg</b>, by decreasing its current level if needed. */
63 void
65  const token_bucket_cfg_t *cfg)
66 {
67  bucket->bucket = MIN(bucket->bucket, cfg->burst);
68 }
69 
70 /**
71  * Given an amount of <b>elapsed</b> time units, and a bucket configuration
72  * <b>cfg</b>, refill the level of <b>bucket</b> accordingly. Note that the
73  * units of time in <b>elapsed</b> must correspond to those used to set the
74  * rate in <b>cfg</b>, or the result will be illogical.
75  */
76 int
78  const token_bucket_cfg_t *cfg,
79  const uint32_t elapsed)
80 {
81  const int was_empty = (bucket->bucket <= 0);
82  /* The casts here prevent an underflow.
83  *
84  * Note that even if the bucket value is negative, subtracting it from
85  * "burst" will still produce a correct result. If this result is
86  * ridiculously high, then the "elapsed > gap / rate" check below
87  * should catch it. */
88  const size_t gap = ((size_t)cfg->burst) - ((size_t)bucket->bucket);
89 
90  if (elapsed > gap / cfg->rate) {
91  bucket->bucket = cfg->burst;
92  } else {
93  bucket->bucket += cfg->rate * elapsed;
94  }
95 
96  return was_empty && bucket->bucket > 0;
97 }
98 
99 /**
100  * Decrement a provided bucket by <b>n</b> units. Note that <b>n</b>
101  * must be nonnegative.
102  */
103 int
105  ssize_t n)
106 {
107  if (BUG(n < 0))
108  return 0;
109  const int becomes_empty = bucket->bucket > 0 && n >= bucket->bucket;
110  bucket->bucket -= n;
111  return becomes_empty;
112 }
113 
114 /** Convert a rate in bytes per second to a rate in bytes per step */
115 STATIC uint32_t
117 {
118  /*
119  The precise calculation we'd want to do is
120 
121  (rate / 1000) * to_approximate_msec(TICKS_PER_STEP). But to minimize
122  rounding error, we do it this way instead, and divide last.
123  */
124  uint64_t units = (uint64_t) rate * TICKS_PER_STEP;
125  uint32_t val = (uint32_t)
127  return val ? val : 1;
128 }
129 
130 /**
131  * Initialize a token bucket in *<b>bucket</b>, set up to allow <b>rate</b>
132  * bytes per second, with a maximum burst of <b>burst</b> bytes. The bucket
133  * is created such that <b>now_ts</b> is the current timestamp. The bucket
134  * starts out full.
135  */
136 void
138  uint32_t rate,
139  uint32_t burst,
140  uint32_t now_ts)
141 {
142  memset(bucket, 0, sizeof(token_bucket_rw_t));
143  token_bucket_rw_adjust(bucket, rate, burst);
144  token_bucket_rw_reset(bucket, now_ts);
145 }
146 
147 /**
148  * Change the configured rate (in bytes per second) and burst (in bytes)
149  * for the token bucket in *<b>bucket</b>.
150  */
151 void
153  uint32_t rate,
154  uint32_t burst)
155 {
156  token_bucket_cfg_init(&bucket->cfg,
158  burst);
159  token_bucket_raw_adjust(&bucket->read_bucket, &bucket->cfg);
160  token_bucket_raw_adjust(&bucket->write_bucket, &bucket->cfg);
161 }
162 
163 /**
164  * Reset <b>bucket</b> to be full, as of timestamp <b>now_ts</b>.
165  */
166 void
168  uint32_t now_ts)
169 {
170  token_bucket_raw_reset(&bucket->read_bucket, &bucket->cfg);
171  token_bucket_raw_reset(&bucket->write_bucket, &bucket->cfg);
172  bucket->last_refilled_at_timestamp = now_ts;
173 }
174 
175 /**
176  * Refill <b>bucket</b> as appropriate, given that the current timestamp
177  * is <b>now_ts</b>.
178  *
179  * Return a bitmask containing TB_READ iff read bucket was empty and became
180  * nonempty, and TB_WRITE iff the write bucket was empty and became nonempty.
181  */
182 int
184  uint32_t now_ts)
185 {
186  const uint32_t elapsed_ticks =
187  (now_ts - bucket->last_refilled_at_timestamp);
188  if (elapsed_ticks > UINT32_MAX-(300*1000)) {
189  /* Either about 48 days have passed since the last refill, or the
190  * monotonic clock has somehow moved backwards. (We're looking at you,
191  * Windows.). We accept up to a 5 minute jump backwards as
192  * "unremarkable".
193  */
194  return 0;
195  }
196  const uint32_t elapsed_steps = elapsed_ticks / TICKS_PER_STEP;
197 
198  if (!elapsed_steps) {
199  /* Note that if less than one whole step elapsed, we don't advance the
200  * time in last_refilled_at. That's intentional: we want to make sure
201  * that we add some bytes to it eventually. */
202  return 0;
203  }
204 
205  int flags = 0;
206  if (token_bucket_raw_refill_steps(&bucket->read_bucket,
207  &bucket->cfg, elapsed_steps))
208  flags |= TB_READ;
209  if (token_bucket_raw_refill_steps(&bucket->write_bucket,
210  &bucket->cfg, elapsed_steps))
211  flags |= TB_WRITE;
212 
213  bucket->last_refilled_at_timestamp = now_ts;
214  return flags;
215 }
216 
217 /**
218  * Decrement the read token bucket in <b>bucket</b> by <b>n</b> bytes.
219  *
220  * Return true if the bucket was nonempty and became empty; return false
221  * otherwise.
222  */
223 int
225  ssize_t n)
226 {
227  return token_bucket_raw_dec(&bucket->read_bucket, n);
228 }
229 
230 /**
231  * Decrement the write token bucket in <b>bucket</b> by <b>n</b> bytes.
232  *
233  * Return true if the bucket was nonempty and became empty; return false
234  * otherwise.
235  */
236 int
238  ssize_t n)
239 {
240  return token_bucket_raw_dec(&bucket->write_bucket, n);
241 }
242 
243 /**
244  * As token_bucket_rw_dec_read and token_bucket_rw_dec_write, in a single
245  * operation. Return a bitmask of TB_READ and TB_WRITE to indicate
246  * which buckets became empty.
247  */
248 int
250  ssize_t n_read, ssize_t n_written)
251 {
252  int flags = 0;
253  if (token_bucket_rw_dec_read(bucket, n_read))
254  flags |= TB_READ;
255  if (token_bucket_rw_dec_write(bucket, n_written))
256  flags |= TB_WRITE;
257  return flags;
258 }
259 
260 /** Initialize a token bucket in <b>bucket</b>, set up to allow <b>rate</b>
261  * per second, with a maximum burst of <b>burst</b>. The bucket is created
262  * such that <b>now_ts</b> is the current timestamp. The bucket starts out
263  * full. */
264 void
266  uint32_t burst, uint32_t now_ts)
267 {
268  memset(bucket, 0, sizeof(token_bucket_ctr_t));
269  token_bucket_ctr_adjust(bucket, rate, burst);
270  token_bucket_ctr_reset(bucket, now_ts);
271 }
272 
273 /** Change the configured rate and burst of the given token bucket object in
274  * <b>bucket</b>. */
275 void
277  uint32_t burst)
278 {
279  token_bucket_cfg_init(&bucket->cfg, rate, burst);
280  token_bucket_raw_adjust(&bucket->counter, &bucket->cfg);
281 }
282 
283 /** Reset <b>bucket</b> to be full, as of timestamp <b>now_ts</b>. */
284 void
286 {
287  token_bucket_raw_reset(&bucket->counter, &bucket->cfg);
288  bucket->last_refilled_at_timestamp = now_ts;
289 }
290 
291 /** Refill <b>bucket</b> as appropriate, given that the current timestamp is
292  * <b>now_ts</b>. */
293 void
295 {
296  const uint32_t elapsed_ticks =
297  (now_ts - bucket->last_refilled_at_timestamp);
298  if (elapsed_ticks > UINT32_MAX-(300*1000)) {
299  /* Either about 48 days have passed since the last refill, or the
300  * monotonic clock has somehow moved backwards. (We're looking at you,
301  * Windows.). We accept up to a 5 minute jump backwards as
302  * "unremarkable".
303  */
304  return;
305  }
306 
307  token_bucket_raw_refill_steps(&bucket->counter, &bucket->cfg,
308  elapsed_ticks);
309  bucket->last_refilled_at_timestamp = now_ts;
310 }
Macro definitions for MIN, MAX, and CLAMP.
uint64_t monotime_coarse_stamp_units_to_approx_msec(uint64_t units)
Definition: compat_time.c:870
Functions and types for monotonic times.
#define STATIC
Definition: testsupport.h:32
int token_bucket_rw_dec(token_bucket_rw_t *bucket, ssize_t n_read, ssize_t n_written)
Definition: token_bucket.c:249
STATIC uint32_t rate_per_sec_to_rate_per_step(uint32_t rate)
Definition: token_bucket.c:116
int token_bucket_rw_dec_read(token_bucket_rw_t *bucket, ssize_t n)
Definition: token_bucket.c:224
void token_bucket_ctr_init(token_bucket_ctr_t *bucket, uint32_t rate, uint32_t burst, uint32_t now_ts)
Definition: token_bucket.c:265
int token_bucket_raw_dec(token_bucket_raw_t *bucket, ssize_t n)
Definition: token_bucket.c:104
void token_bucket_rw_init(token_bucket_rw_t *bucket, uint32_t rate, uint32_t burst, uint32_t now_ts)
Definition: token_bucket.c:137
void token_bucket_raw_adjust(token_bucket_raw_t *bucket, const token_bucket_cfg_t *cfg)
Definition: token_bucket.c:64
void token_bucket_rw_adjust(token_bucket_rw_t *bucket, uint32_t rate, uint32_t burst)
Definition: token_bucket.c:152
void token_bucket_cfg_init(token_bucket_cfg_t *cfg, uint32_t rate, uint32_t burst)
Definition: token_bucket.c:36
void token_bucket_ctr_reset(token_bucket_ctr_t *bucket, uint32_t now_ts)
Definition: token_bucket.c:285
void token_bucket_ctr_refill(token_bucket_ctr_t *bucket, uint32_t now_ts)
Definition: token_bucket.c:294
void token_bucket_ctr_adjust(token_bucket_ctr_t *bucket, uint32_t rate, uint32_t burst)
Definition: token_bucket.c:276
int token_bucket_rw_dec_write(token_bucket_rw_t *bucket, ssize_t n)
Definition: token_bucket.c:237
int token_bucket_raw_refill_steps(token_bucket_raw_t *bucket, const token_bucket_cfg_t *cfg, const uint32_t elapsed)
Definition: token_bucket.c:77
int token_bucket_rw_refill(token_bucket_rw_t *bucket, uint32_t now_ts)
Definition: token_bucket.c:183
void token_bucket_raw_reset(token_bucket_raw_t *bucket, const token_bucket_cfg_t *cfg)
Definition: token_bucket.c:54
void token_bucket_rw_reset(token_bucket_rw_t *bucket, uint32_t now_ts)
Definition: token_bucket.c:167
Headers for token_bucket.c.
#define TOKEN_BUCKET_MAX_BURST
Definition: token_bucket.h:16
Macros to manage assertions, fatal and non-fatal.