Use atomic counters for compressor allocation.

10 files changed, 105 insertions, 10 deletions

diff --git a/src/common/compat_threads.c b/src/common/compat_threads.c
index 7b28a9494a..8d026dd5a7 100644
--- a/src/common/compat_threads.c
+++ b/src/common/compat_threads.c
@@ -352,3 +352,48 @@ alert_sockets_close(alert_sockets_t *socks)
   socks->read_fd = socks->write_fd = -1;
 }
 
+/*
+ * XXXX We might be smart to move to compiler intrinsics or real atomic
+ * XXXX operations at some point.  But not yet.
+ *
+ */
+
+/** Initialize a new atomic counter with the value 0 */
+void
+atomic_counter_init(atomic_counter_t *counter)
+{
+  tor_mutex_init_nonrecursive(&counter->mutex);
+  counter->val = 0;
+}
+/** Clean up all resources held by an atomic counter. */
+void
+atomic_counter_destroy(atomic_counter_t *counter)
+{
+  tor_mutex_uninit(&counter->mutex);
+  memset(counter, 0, sizeof(*counter));
+}
+/** Add a value to an atomic counter. */
+void
+atomic_counter_add(atomic_counter_t *counter, size_t add)
+{
+  tor_mutex_acquire(&counter->mutex);
+  counter->val += add;
+  tor_mutex_release(&counter->mutex);
+}
+/** Subtract a value from an atomic counter. */
+void
+atomic_counter_sub(atomic_counter_t *counter, size_t sub)
+{
+  // this relies on unsigned overflow, but that's fine.
+  atomic_counter_add(counter, -sub);
+}
+/** Return the current value of an atomic counter */
+size_t
+atomic_counter_get(atomic_counter_t *counter)
+{
+  size_t val;
+  tor_mutex_acquire(&counter->mutex);
+  val = counter->val;
+  tor_mutex_release(&counter->mutex);
+  return val;
+}
diff --git a/src/common/compat_threads.h b/src/common/compat_threads.h
index 2943573b23..9fa3d0d0b7 100644
--- a/src/common/compat_threads.h
+++ b/src/common/compat_threads.h
@@ -147,5 +147,19 @@ void *tor_threadlocal_get(tor_threadlocal_t *threadlocal);
  */
 void tor_threadlocal_set(tor_threadlocal_t *threadlocal, void *value);
 
+/**
+ * Atomic counter type; holds a size_t value.
+ */
+typedef struct atomic_counter_t {
+  tor_mutex_t mutex;
+  size_t val;
+} atomic_counter_t;
+
+void atomic_counter_init(atomic_counter_t *counter);
+void atomic_counter_destroy(atomic_counter_t *counter);
+void atomic_counter_add(atomic_counter_t *counter, size_t add);
+void atomic_counter_sub(atomic_counter_t *counter, size_t sub);
+size_t atomic_counter_get(atomic_counter_t *counter);
+
 #endif
 
diff --git a/src/common/compress.c b/src/common/compress.c
index 38b8184573..9332c0017d 100644
--- a/src/common/compress.c
+++ b/src/common/compress.c
@@ -375,3 +375,11 @@ tor_compress_state_size(const tor_compress_state_t *state)
   return 0;
 }
 
+/** Initialize all compression modules. */
+void
+tor_compress_init(void)
+{
+  tor_zlib_init();
+  tor_lzma_init();
+  tor_zstd_init();
+}
diff --git a/src/common/compress.h b/src/common/compress.h
index 2d812e4430..5dd916b961 100644
--- a/src/common/compress.h
+++ b/src/common/compress.h
@@ -79,5 +79,7 @@ void tor_compress_free(tor_compress_state_t *state);
 
 size_t tor_compress_state_size(const tor_compress_state_t *state);
 
+void tor_compress_init(void);
+
 #endif // TOR_COMPRESS_H.
 
diff --git a/src/common/compress_lzma.c b/src/common/compress_lzma.c
index ae0327f581..b0ba4b6ba8 100644
--- a/src/common/compress_lzma.c
+++ b/src/common/compress_lzma.c
@@ -23,7 +23,7 @@
 #endif
 
 /** Total number of bytes allocated for LZMA state. */
-static size_t total_lzma_allocation = 0;
+static atomic_counter_t total_lzma_allocation;
 
 #ifdef HAVE_LZMA
 /** Given <b>level</b> return the memory level. */
@@ -465,6 +465,7 @@ tor_lzma_compress_new(int compress,
     }
   }
 
+  atomic_counter_add(&total_lzma_allocation, result->allocation);
   return result;
 
  err:
@@ -579,7 +580,7 @@ tor_lzma_compress_free(tor_lzma_compress_state_t *state)
   if (state == NULL)
     return;
 
-  total_lzma_allocation -= state->allocation;
+  atomic_counter_sub(&total_lzma_allocation, state->allocation);
 
 #ifdef HAVE_LZMA
   lzma_end(&state->stream);
@@ -600,6 +601,12 @@ tor_lzma_compress_state_size(const tor_lzma_compress_state_t *state)
 size_t
 tor_lzma_get_total_allocation(void)
 {
-  return total_lzma_allocation;
+  return atomic_counter_get(&total_lzma_allocation);
 }
 
+/** Initialize the lzma module */
+void
+tor_lzma_init(void)
+{
+  atomic_counter_init(&total_lzma_allocation);
+}
diff --git a/src/common/compress_lzma.h b/src/common/compress_lzma.h
index 71de56a1c6..60a5f9ce82 100644
--- a/src/common/compress_lzma.h
+++ b/src/common/compress_lzma.h
@@ -47,5 +47,7 @@ size_t tor_lzma_compress_state_size(const tor_lzma_compress_state_t *state);
 
 size_t tor_lzma_get_total_allocation(void);
 
+void tor_lzma_init(void);
+
 #endif // TOR_COMPRESS_LZMA_H.
 
diff --git a/src/common/compress_zlib.c b/src/common/compress_zlib.c
index faa179f7c1..5425f205a0 100644
--- a/src/common/compress_zlib.c
+++ b/src/common/compress_zlib.c
@@ -48,7 +48,7 @@ static size_t tor_zlib_state_size_precalc(int inflate,
                                           int windowbits, int memlevel);
 
 /** Total number of bytes allocated for zlib state */
-static size_t total_zlib_allocation = 0;
+static atomic_counter_t total_zlib_allocation;
 
 /** Given <b>level</b> return the memory level. */
 static int
@@ -437,7 +437,7 @@ tor_zlib_compress_new(int compress_,
   }
   out->allocation = tor_zlib_state_size_precalc(!compress_, bits, memlevel);
 
-  total_zlib_allocation += out->allocation;
+  atomic_counter_add(&total_zlib_allocation, out->allocation);
 
   return out;
 
@@ -519,7 +519,7 @@ tor_zlib_compress_free(tor_zlib_compress_state_t *state)
   if (state == NULL)
     return;
 
-  total_zlib_allocation -= state->allocation;
+  atomic_counter_sub(&total_zlib_allocation, state->allocation);
 
   if (state->compress)
     deflateEnd(&state->stream);
@@ -541,6 +541,12 @@ tor_zlib_compress_state_size(const tor_zlib_compress_state_t *state)
 size_t
 tor_zlib_get_total_allocation(void)
 {
-  return total_zlib_allocation;
+  return atomic_counter_get(&total_zlib_allocation);
 }
 
+/** Set up global state for the zlib module */
+void
+tor_zlib_init(void)
+{
+  atomic_counter_init(&total_zlib_allocation);
+}
diff --git a/src/common/compress_zlib.h b/src/common/compress_zlib.h
index 6e8e5c5136..817ca2785b 100644
--- a/src/common/compress_zlib.h
+++ b/src/common/compress_zlib.h
@@ -47,5 +47,7 @@ size_t tor_zlib_compress_state_size(const tor_zlib_compress_state_t *state);
 
 size_t tor_zlib_get_total_allocation(void);
 
+void tor_zlib_init(void);
+
 #endif // TOR_COMPRESS_ZLIB_H.
 
diff --git a/src/common/compress_zstd.c b/src/common/compress_zstd.c
index 664cce1700..8725c10538 100644
--- a/src/common/compress_zstd.c
+++ b/src/common/compress_zstd.c
@@ -24,7 +24,7 @@
 #endif
 
 /** Total number of bytes allocated for Zstandard state. */
-static size_t total_zstd_allocation = 0;
+static atomic_counter_t total_zstd_allocation;
 
 #ifdef HAVE_ZSTD
 /** Given <b>level</b> return the memory level. */
@@ -446,6 +446,7 @@ tor_zstd_compress_new(int compress,
     }
   }
 
+  atomic_counter_add(&total_zstd_allocation, result->allocation);
   return result;
 
  err:
@@ -578,7 +579,7 @@ tor_zstd_compress_free(tor_zstd_compress_state_t *state)
   if (state == NULL)
     return;
 
-  total_zstd_allocation -= state->allocation;
+  atomic_counter_sub(&total_zstd_allocation, state->allocation);
 
 #ifdef HAVE_ZSTD
   if (state->compress) {
@@ -604,6 +605,12 @@ tor_zstd_compress_state_size(const tor_zstd_compress_state_t *state)
 size_t
 tor_zstd_get_total_allocation(void)
 {
-  return total_zstd_allocation;
+  return atomic_counter_get(&total_zstd_allocation);
 }
 
+/** Initialize the zstd module */
+void
+tor_zstd_init(void)
+{
+  atomic_counter_init(&total_zstd_allocation);
+}
diff --git a/src/common/compress_zstd.h b/src/common/compress_zstd.h
index 663cbdd2ff..186dc3be9f 100644
--- a/src/common/compress_zstd.h
+++ b/src/common/compress_zstd.h
@@ -47,5 +47,7 @@ size_t tor_zstd_compress_state_size(const tor_zstd_compress_state_t *state);
 
 size_t tor_zstd_get_total_allocation(void);
 
+void tor_zstd_init(void);
+
 #endif // TOR_COMPRESS_ZSTD_H.