Make bufferevents work with TokenBucketRefillInterval

5 files changed, 33 insertions, 16 deletions

diff --git a/src/common/compat_libevent.c b/src/common/compat_libevent.c
index beae9502da..3201738701 100644
--- a/src/common/compat_libevent.c
+++ b/src/common/compat_libevent.c
@@ -169,6 +169,7 @@ struct event_base *the_event_base = NULL;
 
 #ifdef USE_BUFFEREVENTS
 static int using_iocp_bufferevents = 0;
+static void tor_libevent_set_tick_timeout(int msec_per_tick);
 
 int
 tor_libevent_using_iocp_bufferevents(void)
@@ -236,6 +237,10 @@ tor_libevent_initialize(tor_libevent_cfg *torcfg)
       "You have a *VERY* old version of libevent.  It is likely to be buggy; "
       "please build Tor with a more recent version.");
 #endif
+
+#ifdef USE_BUFFEREVENTS
+  tor_libevent_set_tick_timeout(torcfg->msec_per_tick);
+#endif
 }
 
 /** Return the current Libevent event base that we're set up to use. */
@@ -598,26 +603,29 @@ static const struct timeval *one_tick = NULL;
 /**
  * Return a special timeout to be passed whenever libevent's O(1) timeout
  * implementation should be used. Only use this when the timer is supposed
- * to fire after 1 / TOR_LIBEVENT_TICKS_PER_SECOND seconds have passed.
+ * to fire after msec_per_tick ticks have elapsed.
 */
 const struct timeval *
 tor_libevent_get_one_tick_timeout(void)
 {
-  if (PREDICT_UNLIKELY(one_tick == NULL)) {
-    struct event_base *base = tor_libevent_get_base();
-    struct timeval tv;
-    if (TOR_LIBEVENT_TICKS_PER_SECOND == 1) {
-      tv.tv_sec = 1;
-      tv.tv_usec = 0;
-    } else {
-      tv.tv_sec = 0;
-      tv.tv_usec = 1000000 / TOR_LIBEVENT_TICKS_PER_SECOND;
-    }
-    one_tick = event_base_init_common_timeout(base, &tv);
-  }
+  tor_assert(one_tick);
   return one_tick;
 }
 
+/** Initialize the common timeout that we'll use to refill the buckets every
+ * time a tick elapses. */
+static void
+tor_libevent_set_tick_timeout(int msec_per_tick)
+{
+  struct event_base *base = tor_libevent_get_base();
+  struct timeval tv;
+
+  tor_assert(! one_tick);
+  tv.tv_sec = msec_per_tick / 1000;
+  tv.tv_usec = (msec_per_tick % 1000) * 1000;
+  one_tick = event_base_init_common_timeout(base, &tv);
+}
+
 static struct bufferevent *
 tor_get_root_bufferevent(struct bufferevent *bev)
 {
diff --git a/src/common/compat_libevent.h b/src/common/compat_libevent.h
index 15b0fc273b..0247297177 100644
--- a/src/common/compat_libevent.h
+++ b/src/common/compat_libevent.h
@@ -62,6 +62,7 @@ int tor_event_base_loopexit(struct event_base *base, struct timeval *tv);
 typedef struct tor_libevent_cfg {
   int disable_iocp;
   int num_cpus;
+  int msec_per_tick;
 } tor_libevent_cfg;
 
 void tor_libevent_initialize(tor_libevent_cfg *cfg);
@@ -73,7 +74,6 @@ void tor_check_libevent_header_compatibility(void);
 const char *tor_libevent_get_version_str(void);
 
 #ifdef USE_BUFFEREVENTS
-#define TOR_LIBEVENT_TICKS_PER_SECOND 3
 const struct timeval *tor_libevent_get_one_tick_timeout(void);
 int tor_libevent_using_iocp_bufferevents(void);
 int tor_set_bufferevent_rate_limit(struct bufferevent *bev,
diff --git a/src/or/config.c b/src/or/config.c
index 5a7a590685..e22f539ade 100644
--- a/src/or/config.c
+++ b/src/or/config.c
@@ -5645,6 +5645,7 @@ init_libevent(const or_options_t *options)
   memset(&cfg, 0, sizeof(cfg));
   cfg.disable_iocp = options->DisableIOCP;
   cfg.num_cpus = get_num_cpus(options);
+  cfg.msec_per_tick = options->TokenBucketRefillInterval;
 
   tor_libevent_initialize(&cfg);
 
diff --git a/src/or/connection.c b/src/or/connection.c
index cb93a81e4e..9dd92972e1 100644
--- a/src/or/connection.c
+++ b/src/or/connection.c
@@ -2561,7 +2561,10 @@ connection_bucket_init(void)
     burst = options->BandwidthBurst;
   }
 
-  rate /= TOR_LIBEVENT_TICKS_PER_SECOND;
+  /* This can't overflow, since TokenBucketRefillInterval <= 1000,
+   * and rate started out less than INT32_MAX. */
+  rate = (rate * options->TokenBucketRefillInterval) / 1000;
+
   bucket_cfg = ev_token_bucket_cfg_new((uint32_t)rate, (uint32_t)burst,
                                        (uint32_t)rate, (uint32_t)burst,
                                        tick);
diff --git a/src/or/connection_or.c b/src/or/connection_or.c
index a75444e1ed..29f0f8de72 100644
--- a/src/or/connection_or.c
+++ b/src/or/connection_or.c
@@ -580,7 +580,12 @@ connection_or_update_token_buckets_helper(or_connection_t *conn, int reset,
   {
     const struct timeval *tick = tor_libevent_get_one_tick_timeout();
     struct ev_token_bucket_cfg *cfg, *old_cfg;
-    int rate_per_tick = rate / TOR_LIBEVENT_TICKS_PER_SECOND;
+    int64_t rate64 = (((int64_t)rate) * options->TokenBucketRefillInterval)
+      / 1000;
+    /* This can't overflow, since TokenBucketRefillInterval <= 1000,
+     * and rate started out less than INT_MAX. */
+    int rate_per_tick = (int) rate64;
+
     cfg = ev_token_bucket_cfg_new(rate_per_tick, burst, rate_per_tick,
                                   burst, tick);
     old_cfg = conn->bucket_cfg;