aboutsummaryrefslogtreecommitdiff
path: root/src/lib/evloop/token_bucket.h
blob: 460dad23e4a99f8adca2caf94a4921e8234eca73 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
/* Copyright (c) 2018-2020, The Tor Project, Inc. */
/* See LICENSE for licensing information */

/**
 * \file token_bucket.h
 * \brief Headers for token_bucket.c
 **/

#ifndef TOR_TOKEN_BUCKET_H
#define TOR_TOKEN_BUCKET_H

#include "lib/cc/torint.h"
#include "lib/testsupport/testsupport.h"

/** Largest allowable burst value for a token buffer. */
#define TOKEN_BUCKET_MAX_BURST INT32_MAX

/** A generic token buffer configuration: determines the number of tokens
 * added to the bucket in each time unit (the "rate"), and the maximum number
 * of tokens in the bucket (the "burst") */
typedef struct token_bucket_cfg_t {
  uint32_t rate;
  int32_t burst;
} token_bucket_cfg_t;

/** A raw token bucket, decoupled from its configuration and timestamp. */
typedef struct token_bucket_raw_t {
  int32_t bucket;
} token_bucket_raw_t;

void token_bucket_cfg_init(token_bucket_cfg_t *cfg,
                           uint32_t rate,
                           uint32_t burst);

void token_bucket_raw_adjust(token_bucket_raw_t *bucket,
                             const token_bucket_cfg_t *cfg);

void token_bucket_raw_reset(token_bucket_raw_t *bucket,
                            const token_bucket_cfg_t *cfg);

int token_bucket_raw_dec(token_bucket_raw_t *bucket,
                         ssize_t n);

int token_bucket_raw_refill_steps(token_bucket_raw_t *bucket,
                                  const token_bucket_cfg_t *cfg,
                                  const uint32_t elapsed_steps);

static inline size_t token_bucket_raw_get(const token_bucket_raw_t *bucket);
/** Return the current number of bytes set in a token bucket. */
static inline size_t
token_bucket_raw_get(const token_bucket_raw_t *bucket)
{
  return bucket->bucket >= 0 ? bucket->bucket : 0;
}

/** A convenience type containing all the pieces needed for a coupled
 * read-bucket and write-bucket that have the same rate limit, and which use
 * "timestamp units" (see compat_time.h) for their time. */
typedef struct token_bucket_rw_t {
  token_bucket_cfg_t cfg;
  token_bucket_raw_t read_bucket;
  token_bucket_raw_t write_bucket;
  uint32_t last_refilled_at_timestamp;
} token_bucket_rw_t;

void token_bucket_rw_init(token_bucket_rw_t *bucket,
                          uint32_t rate,
                          uint32_t burst,
                          uint32_t now_ts);

void token_bucket_rw_adjust(token_bucket_rw_t *bucket,
                            uint32_t rate, uint32_t burst);

void token_bucket_rw_reset(token_bucket_rw_t *bucket,
                           uint32_t now_ts);

#define TB_READ 1
#define TB_WRITE 2

int token_bucket_rw_refill(token_bucket_rw_t *bucket,
                           uint32_t now_ts);

int token_bucket_rw_dec_read(token_bucket_rw_t *bucket,
                             ssize_t n);
int token_bucket_rw_dec_write(token_bucket_rw_t *bucket,
                              ssize_t n);

int token_bucket_rw_dec(token_bucket_rw_t *bucket,
                        ssize_t n_read, ssize_t n_written);

static inline size_t token_bucket_rw_get_read(const token_bucket_rw_t *bucket);
static inline size_t
token_bucket_rw_get_read(const token_bucket_rw_t *bucket)
{
  return token_bucket_raw_get(&bucket->read_bucket);
}

static inline size_t token_bucket_rw_get_write(
                                            const token_bucket_rw_t *bucket);
static inline size_t
token_bucket_rw_get_write(const token_bucket_rw_t *bucket)
{
  return token_bucket_raw_get(&bucket->write_bucket);
}

/**
 * A specialized bucket containing a single counter.
 */

typedef struct token_bucket_ctr_t {
  token_bucket_cfg_t cfg;
  token_bucket_raw_t counter;
  uint32_t last_refilled_at_timestamp;
} token_bucket_ctr_t;

void token_bucket_ctr_init(token_bucket_ctr_t *bucket, uint32_t rate,
                           uint32_t burst, uint32_t now_ts);
void token_bucket_ctr_adjust(token_bucket_ctr_t *bucket, uint32_t rate,
                             uint32_t burst);
void token_bucket_ctr_reset(token_bucket_ctr_t *bucket, uint32_t now_ts);
void token_bucket_ctr_refill(token_bucket_ctr_t *bucket, uint32_t now_ts);

static inline bool
token_bucket_ctr_dec(token_bucket_ctr_t *bucket, ssize_t n)
{
  return token_bucket_raw_dec(&bucket->counter, n);
}

static inline size_t
token_bucket_ctr_get(const token_bucket_ctr_t *bucket)
{
  return token_bucket_raw_get(&bucket->counter);
}

#ifdef TOKEN_BUCKET_PRIVATE

/* To avoid making the rates too small, we consider units of "steps",
 * where a "step" is defined as this many timestamp ticks.  Keep this
 * a power of two if you can. */
#define TICKS_PER_STEP 16

STATIC uint32_t rate_per_sec_to_rate_per_step(uint32_t rate);

#endif /* defined(TOKEN_BUCKET_PRIVATE) */

#endif /* !defined(TOR_TOKEN_BUCKET_H) */