token_bucket.h (4779B)
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_stamp); 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_stamp); 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_stamp); 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_sec); 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_sec); 121 void token_bucket_ctr_refill(token_bucket_ctr_t *bucket, uint32_t now_ts_sec); 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) */