Line data Source code
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 :
21 : #include "lib/evloop/token_bucket.h"
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
36 248 : token_bucket_cfg_init(token_bucket_cfg_t *cfg,
37 : uint32_t rate,
38 : uint32_t burst)
39 : {
40 248 : tor_assert_nonfatal(rate > 0);
41 248 : tor_assert_nonfatal(burst > 0);
42 248 : if (burst > TOKEN_BUCKET_MAX_BURST)
43 : burst = TOKEN_BUCKET_MAX_BURST;
44 :
45 248 : cfg->rate = rate;
46 248 : cfg->burst = burst;
47 248 : }
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
54 251 : token_bucket_raw_reset(token_bucket_raw_t *bucket,
55 : const token_bucket_cfg_t *cfg)
56 : {
57 251 : bucket->bucket = cfg->burst;
58 251 : }
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
64 262 : token_bucket_raw_adjust(token_bucket_raw_t *bucket,
65 : const token_bucket_cfg_t *cfg)
66 : {
67 262 : bucket->bucket = MIN(bucket->bucket, cfg->burst);
68 262 : }
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
77 494 : token_bucket_raw_refill_steps(token_bucket_raw_t *bucket,
78 : const token_bucket_cfg_t *cfg,
79 : const uint32_t elapsed)
80 : {
81 494 : 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 494 : const size_t gap = ((size_t)cfg->burst) - ((size_t)bucket->bucket);
89 :
90 494 : if (elapsed > gap / cfg->rate) {
91 125 : bucket->bucket = cfg->burst;
92 : } else {
93 369 : bucket->bucket += cfg->rate * elapsed;
94 : }
95 :
96 494 : 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
104 386 : token_bucket_raw_dec(token_bucket_raw_t *bucket,
105 : ssize_t n)
106 : {
107 386 : if (BUG(n < 0))
108 0 : return 0;
109 386 : const int becomes_empty = bucket->bucket > 0 && n >= bucket->bucket;
110 386 : bucket->bucket -= n;
111 386 : return becomes_empty;
112 : }
113 :
114 : /** Convert a rate in bytes per second to a rate in bytes per step */
115 : STATIC uint32_t
116 16 : rate_per_sec_to_rate_per_step(uint32_t rate)
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 16 : uint64_t units = (uint64_t) rate * TICKS_PER_STEP;
125 32 : uint32_t val = (uint32_t)
126 16 : (monotime_coarse_stamp_units_to_approx_msec(units) / 1000);
127 16 : 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
137 8 : token_bucket_rw_init(token_bucket_rw_t *bucket,
138 : uint32_t rate,
139 : uint32_t burst,
140 : uint32_t now_ts)
141 : {
142 8 : memset(bucket, 0, sizeof(token_bucket_rw_t));
143 8 : token_bucket_rw_adjust(bucket, rate, burst);
144 8 : token_bucket_rw_reset(bucket, now_ts);
145 8 : }
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
152 14 : token_bucket_rw_adjust(token_bucket_rw_t *bucket,
153 : uint32_t rate,
154 : uint32_t burst)
155 : {
156 14 : token_bucket_cfg_init(&bucket->cfg,
157 : rate_per_sec_to_rate_per_step(rate),
158 : burst);
159 14 : token_bucket_raw_adjust(&bucket->read_bucket, &bucket->cfg);
160 14 : token_bucket_raw_adjust(&bucket->write_bucket, &bucket->cfg);
161 14 : }
162 :
163 : /**
164 : * Reset <b>bucket</b> to be full, as of timestamp <b>now_ts</b>.
165 : */
166 : void
167 10 : token_bucket_rw_reset(token_bucket_rw_t *bucket,
168 : uint32_t now_ts)
169 : {
170 10 : token_bucket_raw_reset(&bucket->read_bucket, &bucket->cfg);
171 10 : token_bucket_raw_reset(&bucket->write_bucket, &bucket->cfg);
172 10 : bucket->last_refilled_at_timestamp = now_ts;
173 10 : }
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
183 11 : token_bucket_rw_refill(token_bucket_rw_t *bucket,
184 : uint32_t now_ts)
185 : {
186 11 : const uint32_t elapsed_ticks =
187 11 : (now_ts - bucket->last_refilled_at_timestamp);
188 11 : 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 10 : const uint32_t elapsed_steps = elapsed_ticks / TICKS_PER_STEP;
197 :
198 10 : 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 9 : int flags = 0;
206 9 : if (token_bucket_raw_refill_steps(&bucket->read_bucket,
207 9 : &bucket->cfg, elapsed_steps))
208 2 : flags |= TB_READ;
209 9 : if (token_bucket_raw_refill_steps(&bucket->write_bucket,
210 : &bucket->cfg, elapsed_steps))
211 0 : flags |= TB_WRITE;
212 :
213 9 : bucket->last_refilled_at_timestamp = now_ts;
214 9 : 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
224 9 : token_bucket_rw_dec_read(token_bucket_rw_t *bucket,
225 : ssize_t n)
226 : {
227 9 : 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
237 0 : token_bucket_rw_dec_write(token_bucket_rw_t *bucket,
238 : ssize_t n)
239 : {
240 0 : 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
249 0 : token_bucket_rw_dec(token_bucket_rw_t *bucket,
250 : ssize_t n_read, ssize_t n_written)
251 : {
252 0 : int flags = 0;
253 0 : if (token_bucket_rw_dec_read(bucket, n_read))
254 0 : flags |= TB_READ;
255 0 : if (token_bucket_rw_dec_write(bucket, n_written))
256 0 : flags |= TB_WRITE;
257 0 : 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
265 230 : token_bucket_ctr_init(token_bucket_ctr_t *bucket, uint32_t rate,
266 : uint32_t burst, uint32_t now_ts)
267 : {
268 230 : memset(bucket, 0, sizeof(token_bucket_ctr_t));
269 230 : token_bucket_ctr_adjust(bucket, rate, burst);
270 230 : token_bucket_ctr_reset(bucket, now_ts);
271 230 : }
272 :
273 : /** Change the configured rate and burst of the given token bucket object in
274 : * <b>bucket</b>. */
275 : void
276 234 : token_bucket_ctr_adjust(token_bucket_ctr_t *bucket, uint32_t rate,
277 : uint32_t burst)
278 : {
279 234 : token_bucket_cfg_init(&bucket->cfg, rate, burst);
280 234 : token_bucket_raw_adjust(&bucket->counter, &bucket->cfg);
281 234 : }
282 :
283 : /** Reset <b>bucket</b> to be full, as of timestamp <b>now_ts</b>. */
284 : void
285 231 : token_bucket_ctr_reset(token_bucket_ctr_t *bucket, uint32_t now_ts)
286 : {
287 231 : token_bucket_raw_reset(&bucket->counter, &bucket->cfg);
288 231 : bucket->last_refilled_at_timestamp = now_ts;
289 231 : }
290 :
291 : /** Refill <b>bucket</b> as appropriate, given that the current timestamp is
292 : * <b>now_ts</b>. */
293 : void
294 476 : token_bucket_ctr_refill(token_bucket_ctr_t *bucket, uint32_t now_ts)
295 : {
296 476 : const uint32_t elapsed_ticks =
297 476 : (now_ts - bucket->last_refilled_at_timestamp);
298 476 : 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 476 : token_bucket_raw_refill_steps(&bucket->counter, &bucket->cfg,
308 : elapsed_ticks);
309 476 : bucket->last_refilled_at_timestamp = now_ts;
310 : }
|
