token_bucket.c (9844B)
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 token_bucket_cfg_init(token_bucket_cfg_t *cfg, 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 54 token_bucket_raw_reset(token_bucket_raw_t *bucket, 55 const token_bucket_cfg_t *cfg) 56 { 57 bucket->bucket = cfg->burst; 58 } 59 60 /** 61 * Adjust a preexisting token bucket to respect the new configuration 62 * <b>cfg</b>, by decreasing its current level if needed. */ 63 void 64 token_bucket_raw_adjust(token_bucket_raw_t *bucket, 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 77 token_bucket_raw_refill_steps(token_bucket_raw_t *bucket, 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 104 token_bucket_raw_dec(token_bucket_raw_t *bucket, 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 * This is used for the 'rw' style (tick based) token buckets but not for 116 * the 'ctr' style buckets which count seconds. */ 117 STATIC uint32_t 118 rate_per_sec_to_rate_per_step(uint32_t rate) 119 { 120 /* 121 The precise calculation we'd want to do is 122 123 (rate / 1000) * to_approximate_msec(TICKS_PER_STEP). But to minimize 124 rounding error, we do it this way instead, and divide last. 125 */ 126 uint64_t units = (uint64_t) rate * TICKS_PER_STEP; 127 uint32_t val = (uint32_t) 128 (monotime_coarse_stamp_units_to_approx_msec(units) / 1000); 129 return val ? val : 1; 130 } 131 132 /** 133 * Initialize a token bucket in *<b>bucket</b>, set up to allow <b>rate</b> 134 * bytes per second, with a maximum burst of <b>burst</b> bytes. The bucket 135 * is created such that <b>now_ts_stamp</b> is the current time in coarse stamp 136 * units. The bucket starts out full. 137 */ 138 void 139 token_bucket_rw_init(token_bucket_rw_t *bucket, 140 uint32_t rate, 141 uint32_t burst, 142 uint32_t now_ts_stamp) 143 { 144 memset(bucket, 0, sizeof(token_bucket_rw_t)); 145 token_bucket_rw_adjust(bucket, rate, burst); 146 token_bucket_rw_reset(bucket, now_ts_stamp); 147 } 148 149 /** 150 * Change the configured rate (in bytes per second) and burst (in bytes) 151 * for the token bucket in *<b>bucket</b>. 152 */ 153 void 154 token_bucket_rw_adjust(token_bucket_rw_t *bucket, 155 uint32_t rate, 156 uint32_t burst) 157 { 158 token_bucket_cfg_init(&bucket->cfg, 159 rate_per_sec_to_rate_per_step(rate), 160 burst); 161 token_bucket_raw_adjust(&bucket->read_bucket, &bucket->cfg); 162 token_bucket_raw_adjust(&bucket->write_bucket, &bucket->cfg); 163 } 164 165 /** 166 * Reset <b>bucket</b> to be full, as of timestamp <b>now_ts_stamp</b>. 167 */ 168 void 169 token_bucket_rw_reset(token_bucket_rw_t *bucket, 170 uint32_t now_ts_stamp) 171 { 172 token_bucket_raw_reset(&bucket->read_bucket, &bucket->cfg); 173 token_bucket_raw_reset(&bucket->write_bucket, &bucket->cfg); 174 bucket->last_refilled_at_timestamp = now_ts_stamp; 175 } 176 177 /** 178 * Refill <b>bucket</b> as appropriate, given that the current timestamp 179 * is <b>now_ts_stamp</b> in coarse timestamp units. 180 * 181 * Return a bitmask containing TB_READ iff read bucket was empty and became 182 * nonempty, and TB_WRITE iff the write bucket was empty and became nonempty. 183 */ 184 int 185 token_bucket_rw_refill(token_bucket_rw_t *bucket, 186 uint32_t now_ts_stamp) 187 { 188 const uint32_t elapsed_ticks = 189 (now_ts_stamp - bucket->last_refilled_at_timestamp); 190 int flags = 0; 191 192 /* Skip over updates that include an overflow or a very large jump. 193 * This can happen for platform specific reasons, such as the old ~48 194 * day windows timer. */ 195 if (elapsed_ticks <= UINT32_MAX/4) { 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 if (token_bucket_raw_refill_steps(&bucket->read_bucket, 206 &bucket->cfg, elapsed_steps)) 207 flags |= TB_READ; 208 if (token_bucket_raw_refill_steps(&bucket->write_bucket, 209 &bucket->cfg, elapsed_steps)) 210 flags |= TB_WRITE; 211 } 212 213 bucket->last_refilled_at_timestamp = now_ts_stamp; 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 224 token_bucket_rw_dec_read(token_bucket_rw_t *bucket, 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 237 token_bucket_rw_dec_write(token_bucket_rw_t *bucket, 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 249 token_bucket_rw_dec(token_bucket_rw_t *bucket, 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_sec</b> is the current timestamp. The bucket starts 263 * out full. Note that these counters use seconds instead of approximate 264 * milliseconds, in order to allow a lower minimum rate than the rw counters. 265 */ 266 void 267 token_bucket_ctr_init(token_bucket_ctr_t *bucket, uint32_t rate, 268 uint32_t burst, uint32_t now_ts_sec) 269 { 270 memset(bucket, 0, sizeof(token_bucket_ctr_t)); 271 token_bucket_ctr_adjust(bucket, rate, burst); 272 token_bucket_ctr_reset(bucket, now_ts_sec); 273 } 274 275 /** Change the configured rate and burst of the given token bucket object in 276 * <b>bucket</b>. */ 277 void 278 token_bucket_ctr_adjust(token_bucket_ctr_t *bucket, uint32_t rate, 279 uint32_t burst) 280 { 281 token_bucket_cfg_init(&bucket->cfg, rate, burst); 282 token_bucket_raw_adjust(&bucket->counter, &bucket->cfg); 283 } 284 285 /** Reset <b>bucket</b> to be full, as of timestamp <b>now_ts_sec</b>. */ 286 void 287 token_bucket_ctr_reset(token_bucket_ctr_t *bucket, uint32_t now_ts_sec) 288 { 289 token_bucket_raw_reset(&bucket->counter, &bucket->cfg); 290 bucket->last_refilled_at_timestamp = now_ts_sec; 291 } 292 293 /** Refill <b>bucket</b> as appropriate, given that the current timestamp is 294 * <b>now_ts_sec</b> in seconds. */ 295 void 296 token_bucket_ctr_refill(token_bucket_ctr_t *bucket, uint32_t now_ts_sec) 297 { 298 const uint32_t elapsed_sec = 299 (now_ts_sec - bucket->last_refilled_at_timestamp); 300 301 /* Are we detecting a rollover or a similar extremely large jump? This 302 * shouldn't generally happen, but if it does for whatever (possibly 303 * platform-specific) reason, ignore it. */ 304 if (elapsed_sec <= UINT32_MAX/4) { 305 token_bucket_raw_refill_steps(&bucket->counter, &bucket->cfg, 306 elapsed_sec); 307 } 308 bucket->last_refilled_at_timestamp = now_ts_sec; 309 }