Refactor conn stats and add unit tests.

3 files changed, 92 insertions, 32 deletions

diff --git a/src/or/rephist.c b/src/or/rephist.c
index a4d885155..e6398194f 100644
--- a/src/or/rephist.c
+++ b/src/or/rephist.c
@@ -2507,19 +2507,25 @@ bidi_map_free(void)
   HT_CLEAR(bidimap, &bidi_map);
 }
 
-/** Stop collecting connection stats in a way that we can re-start doing
- * so in rep_hist_conn_stats_init(). */
+/** Reset counters for conn statistics. */
 void
-rep_hist_conn_stats_term(void)
+rep_hist_reset_conn_stats(time_t now)
 {
+  start_of_conn_stats_interval = now;
   below_threshold = 0;
   mostly_read = 0;
   mostly_written = 0;
   both_read_and_written = 0;
-  start_of_conn_stats_interval = 0;
   bidi_map_free();
 }
 
+/** Stop collecting connection stats in a way that we can re-start doing
+ * so in rep_hist_conn_stats_init(). */
+void
+rep_hist_conn_stats_term(void)
+{
+  rep_hist_reset_conn_stats(0);
+}
 
 /** We read <b>num_read</b> bytes and wrote <b>num_written</b> from/to OR
  * connection <b>conn_id</b> in second <b>when</b>. If this is the first
@@ -2575,46 +2581,60 @@ rep_hist_note_or_conn_bytes(uint64_t conn_id, size_t num_read,
   }
 }
 
-/** Write conn statistics to $DATADIR/stats/conn-stats and return when
- * we would next want to write conn stats. */
+/** Return a newly allocated string containing the connection statistics
+ * until <b>now</b>, or NULL if we're not collecting conn stats. */
+char *
+rep_hist_format_conn_stats(time_t now)
+{
+  char *result, written[ISO_TIME_LEN+1];
+
+  if (!start_of_conn_stats_interval)
+    return NULL; /* Not initialized. */
+
+  format_iso_time(written, now);
+  tor_asprintf(&result, "conn-stats-end %s (%d s)\n"
+               "conn-bi-direct %d,%d,%d,%d\n",
+               written,
+               (unsigned) (now - start_of_conn_stats_interval),
+               below_threshold,
+               mostly_read,
+               mostly_written,
+               both_read_and_written);
+  return result;
+}
+
+/** If 24 hours have passed since the beginning of the current conn stats
+ * period, write conn stats to $DATADIR/stats/conn-stats (possibly
+ * overwriting an existing file) and reset counters.  Return when we would
+ * next want to write conn stats or 0 if we never want to write. */
 time_t
 rep_hist_conn_stats_write(time_t now)
 {
-  char *statsdir = NULL, *filename = NULL;
-  char written[ISO_TIME_LEN+1];
-  open_file_t *open_file = NULL;
-  FILE *out;
+  char *statsdir = NULL, *filename = NULL, *str = NULL;
 
   if (!start_of_conn_stats_interval)
     return 0; /* Not initialized. */
   if (start_of_conn_stats_interval + WRITE_STATS_INTERVAL > now)
     goto done; /* Not ready to write */
 
-  /* write to file */
+  /* Generate history string. */
+  str = rep_hist_format_conn_stats(now);
+
+  /* Reset counters. */
+  rep_hist_reset_conn_stats(now);
+
+  /* Try to write to disk. */
   statsdir = get_datadir_fname("stats");
-  if (check_private_dir(statsdir, CPD_CREATE) < 0)
+  if (check_private_dir(statsdir, CPD_CREATE) < 0) {
+    log_warn(LD_HIST, "Unable to create stats/ directory!");
     goto done;
+  }
   filename = get_datadir_fname2("stats", "conn-stats");
-  out = start_writing_to_stdio_file(filename, OPEN_FLAGS_APPEND,
-                                    0600, &open_file);
-  if (!out)
-    goto done;
-  format_iso_time(written, now);
-  if (fprintf(out, "conn-stats-end %s (%d s)\n", written,
-              (unsigned) (now - start_of_conn_stats_interval)) < 0)
-    goto done;
-
-  if (fprintf(out, "conn-bi-direct %d,%d,%d,%d\n",
-              below_threshold, mostly_read, mostly_written,
-              both_read_and_written) < 0)
-    goto done;
+  if (write_str_to_file(filename, str, 0) < 0)
+    log_warn(LD_HIST, "Unable to write conn stats to disk!");
 
-  finish_writing_to_file(open_file);
-  open_file = NULL;
-  start_of_conn_stats_interval = now;
  done:
-  if (open_file)
-    abort_writing_to_file(open_file);
+  tor_free(str);
   tor_free(filename);
   tor_free(statsdir);
   return start_of_conn_stats_interval + WRITE_STATS_INTERVAL;
diff --git a/src/or/rephist.h b/src/or/rephist.h
index efc8531ac..6b600a2fa 100644
--- a/src/or/rephist.h
+++ b/src/or/rephist.h
@@ -79,6 +79,8 @@ void rep_hist_buffer_stats_term(void);
 void rep_hist_conn_stats_init(time_t now);
 void rep_hist_note_or_conn_bytes(uint64_t conn_id, size_t num_read,
                                  size_t num_written, time_t when);
+void rep_hist_reset_conn_stats(time_t now);
+char *rep_hist_format_conn_stats(time_t now);
 time_t rep_hist_conn_stats_write(time_t now);
 void rep_hist_conn_stats_term(void);
 
diff --git a/src/test/test.c b/src/test/test.c
index 960a6659a..56a7a9a0c 100644
--- a/src/test/test.c
+++ b/src/test/test.c
@@ -1104,7 +1104,8 @@ test_stats(void)
   char *s = NULL;
   int i;
 
-  /* We shouldn't collect exit stats without initializing them. */
+  /* Start with testing exit port statistics; we shouldn't collect exit
+   * stats without initializing them. */
   rep_hist_note_exit_stream_opened(80);
   rep_hist_note_exit_bytes(80, 100, 10000);
   s = rep_hist_format_exit_stats(now + 86400);
@@ -1149,7 +1150,7 @@ test_stats(void)
   test_assert(!s);
 
   /* Re-start stats, add some bytes, reset stats, and see what history we
-   *  get when observing no streams or bytes at all. */
+   * get when observing no streams or bytes at all. */
   rep_hist_exit_stats_init(now);
   rep_hist_note_exit_stream_opened(80);
   rep_hist_note_exit_bytes(80, 100, 10000);
@@ -1159,6 +1160,43 @@ test_stats(void)
              "exit-kibibytes-written other=0\n"
              "exit-kibibytes-read other=0\n"
              "exit-streams-opened other=0\n", s);
+  tor_free(s);
+
+  /* Continue with testing connection statistics; we shouldn't collect
+   * conn stats without initializing them. */
+  rep_hist_note_or_conn_bytes(1, 20, 400, now);
+  s = rep_hist_format_conn_stats(now + 86400);
+  test_assert(!s);
+
+  /* Initialize stats, note bytes, and generate history string. */
+  rep_hist_conn_stats_init(now);
+  rep_hist_note_or_conn_bytes(1, 30000, 400000, now);
+  rep_hist_note_or_conn_bytes(1, 30000, 400000, now + 5);
+  rep_hist_note_or_conn_bytes(2, 400000, 30000, now + 10);
+  rep_hist_note_or_conn_bytes(2, 400000, 30000, now + 15);
+  s = rep_hist_format_conn_stats(now + 86400);
+  test_streq("conn-stats-end 2010-08-12 13:27:30 (86400 s)\n"
+             "conn-bi-direct 0,0,1,0\n", s);
+  tor_free(s);
+
+  /* Stop collecting stats, add some bytes, and ensure we don't generate
+   * a history string. */
+  rep_hist_conn_stats_term();
+  rep_hist_note_or_conn_bytes(2, 400000, 30000, now + 15);
+  s = rep_hist_format_conn_stats(now + 86400);
+  test_assert(!s);
+
+  /* Re-start stats, add some bytes, reset stats, and see what history we
+   * get when observing no streams or bytes at all. */
+  rep_hist_conn_stats_init(now);
+  rep_hist_note_or_conn_bytes(1, 30000, 400000, now);
+  rep_hist_note_or_conn_bytes(1, 30000, 400000, now + 5);
+  rep_hist_note_or_conn_bytes(2, 400000, 30000, now + 10);
+  rep_hist_note_or_conn_bytes(2, 400000, 30000, now + 15);
+  rep_hist_reset_conn_stats(now);
+  s = rep_hist_format_conn_stats(now + 86400);
+  test_streq("conn-stats-end 2010-08-12 13:27:30 (86400 s)\n"
+             "conn-bi-direct 0,0,0,0\n", s);
 
  done:
   tor_free(s);