aboutsummaryrefslogtreecommitdiff
path: root/src/common/token_bucket.c
blob: 3191568ac3c51469b9879711c74f4f5b37cda036 (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
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
/* Copyright (c) 2018, The Tor Project, Inc. */
/* See LICENSE for licensing information */

/**
 * \file token_bucket.c
 * \brief Functions to use and manipulate token buckets, used for
 *    rate-limiting on connections and globally.
 *
 * Tor uses these token buckets to keep track of bandwidth usage, and
 * sometimes other things too.
 *
 * The time units we use internally are based on "timestamp" units -- see
 * monotime_coarse_to_stamp() for a rationale.
 *
 * Token buckets may become negative.
 **/

#define TOKEN_BUCKET_PRIVATE

#include "token_bucket.h"
#include "util_bug.h"

/** Convert a rate in bytes per second to a rate in bytes per step */
static uint32_t
rate_per_sec_to_rate_per_step(uint32_t rate)
{
  /*
    The precise calculation we'd want to do is

    (rate / 1000) * to_approximate_msec(TICKS_PER_STEP).  But to minimize
    rounding error, we do it this way instead, and divide last.
  */
  return (uint32_t)
    monotime_coarse_stamp_units_to_approx_msec(rate*TICKS_PER_STEP)/1000;
}

/**
 * Initialize a token bucket in *<b>bucket</b>, set up to allow <b>rate</b>
 * bytes per second, with a maximum burst of <b>burst</b> bytes. The bucket
 * is created such that <b>now_ts</b> is the current timestamp.  The bucket
 * starts out full.
 */
void
token_bucket_rw_init(token_bucket_rw_t *bucket,
                  uint32_t rate,
                  uint32_t burst,
                  uint32_t now_ts)
{
  memset(bucket, 0, sizeof(token_bucket_rw_t));
  token_bucket_rw_adjust(bucket, rate, burst);
  token_bucket_rw_reset(bucket, now_ts);
}

/**
 * Change the configured rate (in bytes per second) and burst (in bytes)
 * for the token bucket in *<b>bucket</b>.
 */
void
token_bucket_rw_adjust(token_bucket_rw_t *bucket,
                    uint32_t rate,
                    uint32_t burst)
{
  tor_assert_nonfatal(rate > 0);
  tor_assert_nonfatal(burst > 0);
  if (burst > TOKEN_BUCKET_RW_MAX_BURST)
    burst = TOKEN_BUCKET_RW_MAX_BURST;

  bucket->rate = rate_per_sec_to_rate_per_step(rate);
  bucket->burst = burst;
  bucket->read_bucket = MIN(bucket->read_bucket, (int32_t)burst);
  bucket->write_bucket = MIN(bucket->write_bucket, (int32_t)burst);
}

/**
 * Reset <b>bucket</b> to be full, as of timestamp <b>now_ts</b>.
 */
void
token_bucket_rw_reset(token_bucket_rw_t *bucket,
                   uint32_t now_ts)
{
  bucket->read_bucket = bucket->burst;
  bucket->write_bucket = bucket->burst;
  bucket->last_refilled_at_ts = now_ts;
}

/* Helper: see token_bucket_rw_refill */
static int
refill_single_bucket(int32_t *bucketptr,
                     const uint32_t rate,
                     const int32_t burst,
                     const uint32_t elapsed_steps)
{
  const int was_empty = (*bucketptr <= 0);
  /* The casts here prevent an underflow.
   *
   * Note that even if the bucket value is negative, subtracting it from
   * "burst" will still produce a correct result.  If this result is
   * ridiculously high, then the "elapsed_steps > gap / rate" check below
   * should catch it. */
  const size_t gap = ((size_t)burst) - ((size_t)*bucketptr);

  if (elapsed_steps > gap / rate) {
    *bucketptr = burst;
  } else {
    *bucketptr += rate * elapsed_steps;
  }

  return was_empty && *bucketptr > 0;
}

/**
 * Refill <b>bucket</b> as appropriate, given that the current timestamp
 * is <b>now_ts</b>.
 *
 * Return a bitmask containing TB_READ iff read bucket was empty and became
 * nonempty, and TB_WRITE iff the write bucket was empty and became nonempty.
 */
int
token_bucket_rw_refill(token_bucket_rw_t *bucket,
                    uint32_t now_ts)
{
  const uint32_t elapsed_ticks = (now_ts - bucket->last_refilled_at_ts);
  if (elapsed_ticks > UINT32_MAX-(300*1000)) {
    /* Either about 48 days have passed since the last refill, or the
     * monotonic clock has somehow moved backwards. (We're looking at you,
     * Windows.).  We accept up to a 5 minute jump backwards as
     * "unremarkable".
     */
    return 0;
  }
  const uint32_t elapsed_steps = elapsed_ticks / TICKS_PER_STEP;

  if (!elapsed_steps) {
    /* Note that if less than one whole step elapsed, we don't advance the
     * time in last_refilled_at_ts. That's intentional: we want to make sure
     * that we add some bytes to it eventually. */
    return 0;
  }

  int flags = 0;
  if (refill_single_bucket(&bucket->read_bucket,
                           bucket->rate, bucket->burst, elapsed_steps))
    flags |= TB_READ;
  if (refill_single_bucket(&bucket->write_bucket,
                           bucket->rate, bucket->burst, elapsed_steps))
    flags |= TB_WRITE;

  bucket->last_refilled_at_ts = now_ts;
  return flags;
}

static int
decrement_single_bucket(int32_t *bucketptr,
                        ssize_t n)
{
  if (BUG(n < 0))
    return 0;
  const int becomes_empty = *bucketptr > 0 && n >= *bucketptr;
  *bucketptr -= n;
  return becomes_empty;
}

/**
 * Decrement the read token bucket in <b>bucket</b> by <b>n</b> bytes.
 *
 * Return true if the bucket was nonempty and became empty; return false
 * otherwise.
 */
int
token_bucket_rw_dec_read(token_bucket_rw_t *bucket,
                      ssize_t n)
{
  return decrement_single_bucket(&bucket->read_bucket, n);
}

/**
 * Decrement the write token bucket in <b>bucket</b> by <b>n</b> bytes.
 *
 * Return true if the bucket was nonempty and became empty; return false
 * otherwise.
 */
int
token_bucket_rw_dec_write(token_bucket_rw_t *bucket,
                       ssize_t n)
{
  return decrement_single_bucket(&bucket->write_bucket, n);
}

/**
 * As token_bucket_rw_dec_read and token_bucket_rw_dec_write, in a single
 * operation.
 */
void
token_bucket_rw_dec(token_bucket_rw_t *bucket,
                 ssize_t n_read, ssize_t n_written)
{
  token_bucket_rw_dec_read(bucket, n_read);
  token_bucket_rw_dec_read(bucket, n_written);
}