diff --git a/src/or/connection.c b/src/or/connection.c
index e2fd196dd4..4013e0538f 100644
--- a/src/or/connection.c
+++ b/src/or/connection.c
@@ -119,8 +119,6 @@ static connection_t *connection_listener_new(
static void connection_init(time_t now, connection_t *conn, int type,
int socket_family);
static int connection_handle_listener_read(connection_t *conn, int new_type);
-static int connection_bucket_should_increase(int bucket,
- or_connection_t *conn);
static int connection_finished_flushing(connection_t *conn);
static int connection_flushed_some(connection_t *conn);
static int connection_finished_connecting(connection_t *conn);
@@ -2887,7 +2885,7 @@ connection_bucket_read_limit(connection_t *conn, time_t now)
if (connection_speaks_cells(conn)) {
or_connection_t *or_conn = TO_OR_CONN(conn);
if (conn->state == OR_CONN_STATE_OPEN)
- conn_bucket = or_conn->read_bucket;
+ conn_bucket = token_bucket_get_read(&or_conn->bucket);
base = get_cell_network_size(or_conn->wide_circ_ids);
}
@@ -2919,13 +2917,10 @@ connection_bucket_write_limit(connection_t *conn, time_t now)
}
if (connection_speaks_cells(conn)) {
- /* use the per-conn write limit if it's lower, but if it's less
- * than zero just use zero */
+ /* use the per-conn write limit if it's lower */
or_connection_t *or_conn = TO_OR_CONN(conn);
if (conn->state == OR_CONN_STATE_OPEN)
- if (or_conn->write_bucket < conn_bucket)
- conn_bucket = or_conn->write_bucket >= 0 ?
- or_conn->write_bucket : 0;
+ conn_bucket = MIN(conn_bucket, token_bucket_get_write(&or_conn->bucket));
base = get_cell_network_size(or_conn->wide_circ_ids);
}
@@ -3050,8 +3045,8 @@ connection_buckets_decrement(connection_t *conn, time_t now,
global_read_bucket -= (int)num_read;
global_write_bucket -= (int)num_written;
if (connection_speaks_cells(conn) && conn->state == OR_CONN_STATE_OPEN) {
- TO_OR_CONN(conn)->read_bucket -= (int)num_read;
- TO_OR_CONN(conn)->write_bucket -= (int)num_written;
+ or_connection_t *or_conn = TO_OR_CONN(conn);
+ token_bucket_dec(&or_conn->bucket, num_read, num_written);
}
}
@@ -3072,7 +3067,7 @@ connection_consider_empty_read_buckets(connection_t *conn)
reason = "global relayed read bucket exhausted. Pausing.";
} else if (connection_speaks_cells(conn) &&
conn->state == OR_CONN_STATE_OPEN &&
- TO_OR_CONN(conn)->read_bucket <= 0) {
+ token_bucket_get_read(&TO_OR_CONN(conn)->bucket) <= 0) {
reason = "connection read bucket exhausted. Pausing.";
} else
return; /* all good, no need to stop it */
@@ -3099,7 +3094,7 @@ connection_consider_empty_write_buckets(connection_t *conn)
reason = "global relayed write bucket exhausted. Pausing.";
} else if (connection_speaks_cells(conn) &&
conn->state == OR_CONN_STATE_OPEN &&
- TO_OR_CONN(conn)->write_bucket <= 0) {
+ token_bucket_get_write(&TO_OR_CONN(conn)->bucket) <= 0) {
reason = "connection write bucket exhausted. Pausing.";
} else
return; /* all good, no need to stop it */
@@ -3157,7 +3152,7 @@ connection_bucket_refill_helper(int *bucket, int rate, int burst,
/** Time has passed; increment buckets appropriately. */
void
-connection_bucket_refill(int milliseconds_elapsed, time_t now)
+connection_bucket_refill(int milliseconds_elapsed, time_t now, uint32_t now_ts)
{
const or_options_t *options = get_options();
smartlist_t *conns = get_connection_array();
@@ -3201,22 +3196,9 @@ connection_bucket_refill(int milliseconds_elapsed, time_t now)
SMARTLIST_FOREACH_BEGIN(conns, connection_t *, conn) {
if (connection_speaks_cells(conn)) {
or_connection_t *or_conn = TO_OR_CONN(conn);
- int orbandwidthrate = or_conn->bandwidthrate;
- int orbandwidthburst = or_conn->bandwidthburst;
- if (connection_bucket_should_increase(or_conn->read_bucket, or_conn)) {
- connection_bucket_refill_helper(&or_conn->read_bucket,
- orbandwidthrate,
- orbandwidthburst,
- milliseconds_elapsed,
- "or_conn->read_bucket");
- }
- if (connection_bucket_should_increase(or_conn->write_bucket, or_conn)) {
- connection_bucket_refill_helper(&or_conn->write_bucket,
- orbandwidthrate,
- orbandwidthburst,
- milliseconds_elapsed,
- "or_conn->write_bucket");
+ if (conn->state == OR_CONN_STATE_OPEN) {
+ token_bucket_refill(&or_conn->bucket, now_ts);
}
}
@@ -3226,7 +3208,7 @@ connection_bucket_refill(int milliseconds_elapsed, time_t now)
global_relayed_read_bucket > 0) /* even if we're relayed traffic */
&& (!connection_speaks_cells(conn) ||
conn->state != OR_CONN_STATE_OPEN ||
- TO_OR_CONN(conn)->read_bucket > 0)) {
+ token_bucket_get_read(&TO_OR_CONN(conn)->bucket) > 0)) {
/* and either a non-cell conn or a cell conn with non-empty bucket */
LOG_FN_CONN(conn, (LOG_DEBUG,LD_NET,
"waking up conn (fd %d) for read", (int)conn->s));
@@ -3240,7 +3222,7 @@ connection_bucket_refill(int milliseconds_elapsed, time_t now)
global_relayed_write_bucket > 0) /* even if it's relayed traffic */
&& (!connection_speaks_cells(conn) ||
conn->state != OR_CONN_STATE_OPEN ||
- TO_OR_CONN(conn)->write_bucket > 0)) {
+ token_bucket_get_write(&TO_OR_CONN(conn)->bucket) > 0)) {
LOG_FN_CONN(conn, (LOG_DEBUG,LD_NET,
"waking up conn (fd %d) for write", (int)conn->s));
conn->write_blocked_on_bw = 0;
@@ -3249,22 +3231,6 @@ connection_bucket_refill(int milliseconds_elapsed, time_t now)
} SMARTLIST_FOREACH_END(conn);
}
-/** Is the bucket for connection conn low enough that we
- * should add another pile of tokens to it?
- */
-static int
-connection_bucket_should_increase(int bucket, or_connection_t *conn)
-{
- tor_assert(conn);
-
- if (conn->base_.state != OR_CONN_STATE_OPEN)
- return 0; /* only open connections play the rate limiting game */
- if (bucket >= conn->bandwidthburst)
- return 0;
-
- return 1;
-}
-
/** Read bytes from conn-\>s and process them.
*
* It calls connection_buf_read_from_socket() to bring in any new bytes,
diff --git a/src/or/connection.h b/src/or/connection.h
index 930e5f81ae..acc63640d1 100644
--- a/src/or/connection.h
+++ b/src/or/connection.h
@@ -122,7 +122,8 @@ void connection_mark_all_noncontrol_connections(void);
ssize_t connection_bucket_write_limit(connection_t *conn, time_t now);
int global_write_bucket_low(connection_t *conn, size_t attempt, int priority);
void connection_bucket_init(void);
-void connection_bucket_refill(int seconds_elapsed, time_t now);
+void connection_bucket_refill(int seconds_elapsed, time_t now,
+ uint32_t now_ts);
int connection_handle_read(connection_t *conn);
diff --git a/src/or/connection_or.c b/src/or/connection_or.c
index 267463312c..3afdfa6b5a 100644
--- a/src/or/connection_or.c
+++ b/src/or/connection_or.c
@@ -793,18 +793,10 @@ connection_or_update_token_buckets_helper(or_connection_t *conn, int reset,
(int)options->BandwidthBurst, 1, INT32_MAX);
}
- conn->bandwidthrate = rate;
- conn->bandwidthburst = burst;
- if (reset) { /* set up the token buckets to be full */
- conn->read_bucket = conn->write_bucket = burst;
- return;
+ token_bucket_adjust(&conn->bucket, rate, burst);
+ if (reset) {
+ token_bucket_reset(&conn->bucket, monotime_coarse_get_stamp());
}
- /* If the new token bucket is smaller, take out the extra tokens.
- * (If it's larger, don't -- the buckets can grow to reach the cap.) */
- if (conn->read_bucket > burst)
- conn->read_bucket = burst;
- if (conn->write_bucket > burst)
- conn->write_bucket = burst;
}
/** Either our set of relays or our per-conn rate limits have changed.
diff --git a/src/or/main.c b/src/or/main.c
index d1df11af50..b12effce1f 100644
--- a/src/or/main.c
+++ b/src/or/main.c
@@ -2426,8 +2426,10 @@ refill_callback(periodic_timer_t *timer, void *arg)
if (accounting_is_enabled(options) && milliseconds_elapsed >= 0)
accounting_add_bytes(bytes_read, bytes_written, seconds_rolled_over);
- if (milliseconds_elapsed > 0)
- connection_bucket_refill(milliseconds_elapsed, (time_t)now.tv_sec);
+ if (milliseconds_elapsed > 0) {
+ connection_bucket_refill(milliseconds_elapsed, (time_t)now.tv_sec,
+ monotime_coarse_get_stamp());
+ }
stats_prev_global_read_bucket = global_read_bucket;
stats_prev_global_write_bucket = global_write_bucket;
diff --git a/src/or/or.h b/src/or/or.h
index bcce33755d..a826cacbf3 100644
--- a/src/or/or.h
+++ b/src/or/or.h
@@ -80,6 +80,7 @@
#include "crypto_curve25519.h"
#include "crypto_ed25519.h"
#include "tor_queue.h"
+#include "token_bucket.h"
#include "util_format.h"
#include "hs_circuitmap.h"
@@ -1652,13 +1653,8 @@ typedef struct or_connection_t {
time_t timestamp_lastempty; /**< When was the outbuf last completely empty?*/
- /* bandwidth* and *_bucket only used by ORs in OPEN state: */
- int bandwidthrate; /**< Bytes/s added to the bucket. (OPEN ORs only.) */
- int bandwidthburst; /**< Max bucket size for this conn. (OPEN ORs only.) */
- int read_bucket; /**< When this hits 0, stop receiving. Every second we
- * add 'bandwidthrate' to this, capping it at
- * bandwidthburst. (OPEN ORs only) */
- int write_bucket; /**< When this hits 0, stop writing. Like read_bucket. */
+ token_bucket_t bucket; /**< Used for rate limiting when the connection is
+ * in state CONN_OPEN. */
/*
* Count the number of bytes flushed out on this orconn, and the number of