New torrc option to allow bucket refill intervals of less than 1 sec

Implements bug3630.

1 files changed, 86 insertions, 24 deletions

diff --git a/src/or/main.c b/src/or/main.c
index ad5558e649..1a95e31c04 100644
--- a/src/or/main.c
+++ b/src/or/main.c
@@ -91,10 +91,10 @@ static int stats_prev_global_read_bucket;
 /** What was the write bucket before the last second_elapsed_callback() call?
  * (used to determine how many bytes we've written). */
 static int stats_prev_global_write_bucket;
-#else
+#endif
+
 static uint64_t stats_prev_n_read = 0;
 static uint64_t stats_prev_n_written = 0;
-#endif
 
 /* XXX we might want to keep stats about global_relayed_*_bucket too. Or not.*/
 /** How many bytes have we read since we started the process? */
@@ -1507,9 +1507,6 @@ second_elapsed_callback(periodic_timer_t *timer, void *arg)
   size_t bytes_written;
   size_t bytes_read;
   int seconds_elapsed;
-#ifdef USE_BUFFEREVENTS
-  uint64_t cur_read,cur_written;
-#endif
   const or_options_t *options = get_options();
   (void)timer;
   (void)arg;
@@ -1523,30 +1520,28 @@ second_elapsed_callback(periodic_timer_t *timer, void *arg)
   /* the second has rolled over. check more stuff. */
   seconds_elapsed = current_second ? (int)(now - current_second) : 0;
 #ifdef USE_BUFFEREVENTS
-  connection_get_rate_limit_totals(&cur_read, &cur_written);
-  bytes_written = (size_t)(cur_written - stats_prev_n_written);
-  bytes_read = (size_t)(cur_read - stats_prev_n_read);
+  {
+    uint64_t cur_read,cur_written;
+    connection_get_rate_limit_totals(&cur_read, &cur_written);
+    bytes_written = (size_t)(cur_written - stats_prev_n_written);
+    bytes_read = (size_t)(cur_read - stats_prev_n_read);
+    stats_n_bytes_read += bytes_read;
+    stats_n_bytes_written += bytes_written;
+    if (accounting_is_enabled(options) && seconds_elapsed >= 0)
+      accounting_add_bytes(bytes_read, bytes_written, seconds_elapsed);
+    stats_prev_n_written = cur_written;
+    stats_prev_n_read = cur_read;
+  }
 #else
-  bytes_written = stats_prev_global_write_bucket - global_write_bucket;
-  bytes_read = stats_prev_global_read_bucket - global_read_bucket;
+  bytes_read = (size_t)(stats_n_bytes_read - stats_prev_n_read);
+  bytes_written = (size_t)(stats_n_bytes_written - stats_prev_n_written);
+  stats_prev_n_read = stats_n_bytes_read;
+  stats_prev_n_written = stats_n_bytes_written;
 #endif
-  stats_n_bytes_read += bytes_read;
-  stats_n_bytes_written += bytes_written;
-  if (accounting_is_enabled(options) && seconds_elapsed >= 0)
-    accounting_add_bytes(bytes_read, bytes_written, seconds_elapsed);
+
   control_event_bandwidth_used((uint32_t)bytes_read,(uint32_t)bytes_written);
   control_event_stream_bandwidth_used();
 
-  if (seconds_elapsed > 0)
-    connection_bucket_refill(seconds_elapsed, now);
-#ifdef USE_BUFFEREVENTS
-  stats_prev_n_written = cur_written;
-  stats_prev_n_read = cur_read;
-#else
-  stats_prev_global_read_bucket = global_read_bucket;
-  stats_prev_global_write_bucket = global_write_bucket;
-#endif
-
   if (server_mode(options) &&
       !we_are_hibernating() &&
       seconds_elapsed > 0 &&
@@ -1594,6 +1589,57 @@ second_elapsed_callback(periodic_timer_t *timer, void *arg)
   current_second = now; /* remember which second it is, for next time */
 }
 
+#ifndef USE_BUFFEREVENTS
+/** Timer: used to invoke refill_callback(). */
+static periodic_timer_t *refill_timer = NULL;
+
+/** Libevent callback: invoked periodically to refill token buckets
+ * and count r/w bytes. It is only used when bufferevents are disabled. */
+static void
+refill_callback(periodic_timer_t *timer, void *arg)
+{
+  static struct timeval current_millisecond;
+  struct timeval now;
+
+  size_t bytes_written;
+  size_t bytes_read;
+  int milliseconds_elapsed = 0;
+  int seconds_rolled_over = 0;
+
+  const or_options_t *options = get_options();
+
+  (void)timer;
+  (void)arg;
+
+  tor_gettimeofday(&now);
+
+  /* If this is our first time, no time has passed. */
+  if (current_millisecond.tv_sec) {
+    long mdiff = tv_mdiff(&current_millisecond, &now);
+    if (mdiff > INT_MAX)
+      mdiff = INT_MAX;
+    milliseconds_elapsed = (int)mdiff;
+    seconds_rolled_over = (int)(now.tv_sec - current_millisecond.tv_sec);
+  }
+
+  bytes_written = stats_prev_global_write_bucket - global_write_bucket;
+  bytes_read = stats_prev_global_read_bucket - global_read_bucket;
+
+  stats_n_bytes_read += bytes_read;
+  stats_n_bytes_written += bytes_written;
+  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, now.tv_sec);
+
+  stats_prev_global_read_bucket = global_read_bucket;
+  stats_prev_global_write_bucket = global_write_bucket;
+
+  current_millisecond = now; /* remember what time it is, for next time */
+}
+#endif
+
 #ifndef MS_WINDOWS
 /** Called when a possibly ignorable libevent error occurs; ensures that we
  * don't get into an infinite loop by ignoring too many errors from
@@ -1791,6 +1837,22 @@ do_main_loop(void)
     tor_assert(second_timer);
   }
 
+#ifndef USE_BUFFEREVENTS
+  if (!refill_timer) {
+    struct timeval refill_interval;
+    int msecs = get_options()->TokenBucketRefillInterval;
+
+    refill_interval.tv_sec =  msecs/1000;
+    refill_interval.tv_usec = (msecs%1000)*1000;
+
+    refill_timer = periodic_timer_new(tor_libevent_get_base(),
+                                      &refill_interval,
+                                      refill_callback,
+                                      NULL);
+    tor_assert(refill_timer);
+  }
+#endif
+
   for (;;) {
     if (nt_service_is_stopping())
       return 0;