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.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        1273 : token_bucket_raw_get(const token_bucket_raw_t *bucket)
      52             : {
      53        1273 :   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           0 : token_bucket_rw_get_read(const token_bucket_rw_t *bucket)
      94             : {
      95           0 :   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          70 : token_bucket_rw_get_write(const token_bucket_rw_t *bucket)
     102             : {
     103          70 :   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         377 : token_bucket_ctr_dec(token_bucket_ctr_t *bucket, ssize_t n)
     125             : {
     126         377 :   return token_bucket_raw_dec(&bucket->counter, n);
     127             : }
     128             : 
     129             : static inline size_t
     130        1203 : token_bucket_ctr_get(const token_bucket_ctr_t *bucket)
     131             : {
     132        1203 :   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) */