perf(compression,tcp): reuse zstd contexts; enable TCP_NODELAY

src/shared/compression.c

@@ -2,11 +2,12 @@
 #include "data.h"
 #include "log.h"
 #include "protocol.h"
-#include <stdlib.h>
 #include <limits.h>
 #include <stdint.h>
+#include <stdlib.h>
 #include <string.h>
 #include <strings.h>
+#include <threads.h>
 #include <unistd.h>
 #include <zstd.h>
 
@@ -39,6 +40,96 @@ bool compression_should_skip_with_suffixes(const char* path, char* const* suffix
   return false;
 }
 
+/* Per-thread cache of zstd contexts plus the grow-only compression scratch
+ * buffer.  zstd contexts are stateful and not safe to share between threads,
+ * so each thread keeps its own (see compression_get_thread_ctx).  The cache is
+ * stored in a C11 thread-specific storage slot whose destructor releases the
+ * contexts when the thread exits; this keeps LeakSanitizer clean for the
+ * short-lived sender/receiver/scanner worker threads without every worker
+ * entry point having to remember to call compression_free_thread_contexts().
+ * The main thread's slot is not torn down by tss at process exit, so an atexit
+ * hook releases it (and compression_free_thread_contexts allows eager
+ * release). */
+typedef struct {
+  ZSTD_CCtx* cctx;
+  ZSTD_DCtx* dctx;
+  void* out_buf;  /* reusable ZSTD_compressBound-sized output scratch */
+  size_t out_cap; /* bytes currently allocated for out_buf */
+  int level;      /* compression level currently applied to cctx */
+  int workers;    /* nbWorkers currently applied to cctx */
+  bool params_set;
+  bool cached; /* false when the TSS slot could not be used: caller owns */
+} CompressionThreadCtx;
+
+static once_flag compression_tls_once = ONCE_FLAG_INIT;
+static tss_t compression_tls_key;
+static bool compression_tls_ready;
+
+static void compression_tls_make_key(void);
+
+static void compression_ctx_free(CompressionThreadCtx* ctx) {
+  if (!ctx)
+    return;
+  if (ctx->cctx)
+    ZSTD_freeCCtx(ctx->cctx);
+  if (ctx->dctx)
+    ZSTD_freeDCtx(ctx->dctx);
+  free(ctx->out_buf);
+  free(ctx);
+}
+
+static void compression_tls_destructor(void* value) {
+  compression_ctx_free((CompressionThreadCtx*)value);
+}
+
+void compression_free_thread_contexts(void) {
+  call_once(&compression_tls_once, compression_tls_make_key);
+  if (!compression_tls_ready)
+    return;
+  CompressionThreadCtx* ctx = (CompressionThreadCtx*)tss_get(compression_tls_key);
+  if (!ctx)
+    return;
+  /* Clear the slot first so the thread-exit destructor cannot free it twice. */
+  tss_set(compression_tls_key, NULL);
+  compression_ctx_free(ctx);
+}
+
+static void compression_atexit_cleanup(void) {
+  compression_free_thread_contexts();
+}
+
+static void compression_tls_make_key(void) {
+  if (tss_create(&compression_tls_key, compression_tls_destructor) == thrd_success) {
+    compression_tls_ready = true;
+    atexit(compression_atexit_cleanup);
+  }
+}
+
+static CompressionThreadCtx* compression_get_thread_ctx(void) {
+  call_once(&compression_tls_once, compression_tls_make_key);
+  if (!compression_tls_ready) {
+    /* Extremely unlikely: fall back to an uncached context the caller frees. */
+    return (CompressionThreadCtx*)calloc(1, sizeof(CompressionThreadCtx));
+  }
+  CompressionThreadCtx* ctx = (CompressionThreadCtx*)tss_get(compression_tls_key);
+  if (ctx)
+    return ctx;
+  ctx = (CompressionThreadCtx*)calloc(1, sizeof(CompressionThreadCtx));
+  if (!ctx)
+    return NULL;
+  ctx->cached = true;
+  if (tss_set(compression_tls_key, ctx) != thrd_success)
+    ctx->cached = false;
+  return ctx;
+}
+
+/* Release an uncached context immediately; cached contexts are owned by the
+ * thread's TSS slot and freed on thread exit / compression_free_thread_contexts. */
+static void compression_ctx_put(CompressionThreadCtx* ctx) {
+  if (ctx && !ctx->cached)
+    compression_ctx_free(ctx);
+}
+
 Data* data_compress(Data* data_to_compress, int compression_level) {
   return data_compress_with_threads(data_to_compress, compression_level, 0);
 }
@@ -50,68 +141,102 @@ Data* data_compress_with_threads(Data* data_to_compress, int compression_level,
     return NULL;
   log_message(LOG_LEVEL_DEBUG, "Starting to compress data");
   size_t dst_size = ZSTD_compressBound(data_to_compress->size);
-  Data* compressed_data = data_create_empty(dst_size);
-  if (compressed_data == NULL)
-    return NULL;
 
-  ZSTD_CCtx* cctx = ZSTD_createCCtx();
-  if (!cctx) {
-    log_message(LOG_LEVEL_ERROR, "Failed to create ZSTD compression context");
-    data_destroy(compressed_data);
+  CompressionThreadCtx* ctx = compression_get_thread_ctx();
+  if (ctx == NULL) {
+    log_message(LOG_LEVEL_ERROR, "Failed to allocate ZSTD compression context");
     return NULL;
   }
+  Data* compressed_data = NULL;
 
-  size_t zret = ZSTD_CCtx_setParameter(cctx, ZSTD_c_compressionLevel, compression_level);
-  if (ZSTD_isError(zret)) {
-    log_message(LOG_LEVEL_ERROR, "Failed to set compression level: %s", ZSTD_getErrorName(zret));
-    ZSTD_freeCCtx(cctx);
-    data_destroy(compressed_data);
-    return NULL;
+  if (!ctx->cctx) {
+    ctx->cctx = ZSTD_createCCtx();
+    if (!ctx->cctx) {
+      log_message(LOG_LEVEL_ERROR, "Failed to create ZSTD compression context");
+      goto cleanup;
+    }
+    ctx->params_set = false;
   }
 
+  /* Reset only the session: parameters (and any already-allocated zstd worker
+   * pool) stay attached to the context, so compressing the next file does not
+   * rebuild the pool. */
+  ZSTD_CCtx_reset(ctx->cctx, ZSTD_reset_session_only);
+
+  if (!ctx->params_set || ctx->level != compression_level) {
+    size_t zret = ZSTD_CCtx_setParameter(ctx->cctx, ZSTD_c_compressionLevel, compression_level);
+    if (ZSTD_isError(zret)) {
+      log_message(LOG_LEVEL_ERROR, "Failed to set compression level: %s", ZSTD_getErrorName(zret));
+      goto cleanup;
+    }
+    ctx->level = compression_level;
+  }
+
+  int available_threads = 0;
   if (compression_threads > 0) {
     long online_cpus = sysconf(_SC_NPROCESSORS_ONLN);
-    int available_threads = online_cpus > 0 && online_cpus < compression_threads
-                                ? (int)online_cpus
-                                : compression_threads;
-    zret = ZSTD_CCtx_setParameter(cctx, ZSTD_c_nbWorkers, available_threads);
+    available_threads = online_cpus > 0 && online_cpus < compression_threads ? (int)online_cpus
+                                                                             : compression_threads;
+  }
+  if (!ctx->params_set || ctx->workers != available_threads) {
+    size_t zret = ZSTD_CCtx_setParameter(ctx->cctx, ZSTD_c_nbWorkers, available_threads);
     if (ZSTD_isError(zret)) {
       log_message(LOG_LEVEL_ERROR, "Failed to set compression threads: %s",
                   ZSTD_getErrorName(zret));
-      ZSTD_freeCCtx(cctx);
-      data_destroy(compressed_data);
-      return NULL;
+      goto cleanup;
     }
+    ctx->workers = available_threads;
+  }
+  ctx->params_set = true;
+
+  if (available_threads > 0) {
     /* Streaming compression needs the source size before threaded mode can end a frame. */
-    zret = ZSTD_CCtx_setPledgedSrcSize(cctx, data_to_compress->size);
+    size_t zret = ZSTD_CCtx_setPledgedSrcSize(ctx->cctx, data_to_compress->size);
     if (ZSTD_isError(zret)) {
       log_message(LOG_LEVEL_ERROR, "Failed to set compression source size: %s",
                   ZSTD_getErrorName(zret));
-      ZSTD_freeCCtx(cctx);
-      data_destroy(compressed_data);
-      return NULL;
+      goto cleanup;
     }
   }
 
+  if (ctx->out_cap < dst_size) {
+    void* grown = protocol_realloc(ctx->out_buf, dst_size);
+    if (grown == NULL) {
+      log_message(LOG_LEVEL_ERROR, "Failed to allocate compression buffer");
+      goto cleanup;
+    }
+    ctx->out_buf = grown;
+    ctx->out_cap = dst_size;
+  }
+
   ZSTD_inBuffer input = {data_to_compress->data, data_to_compress->size, 0};
-  ZSTD_outBuffer output = {compressed_data->data, dst_size, 0};
+  ZSTD_outBuffer output = {ctx->out_buf, dst_size, 0};
 
   size_t ret;
   do {
-    ret = ZSTD_compressStream2(cctx, &output, &input, ZSTD_e_end);
+    ret = ZSTD_compressStream2(ctx->cctx, &output, &input, ZSTD_e_end);
     if (ZSTD_isError(ret)) {
       log_message(LOG_LEVEL_ERROR, "Compression failed: %s", ZSTD_getErrorName(ret));
-      ZSTD_freeCCtx(cctx);
-      data_destroy(compressed_data);
-      return NULL;
+      goto cleanup;
     }
   } while (ret > 0);
 
+  /* Hand off an exactly-sized copy; the scratch buffer stays cached so the next
+   * call does not reallocate a ZSTD_compressBound-sized block. */
+  compressed_data = data_create_empty(output.pos);
+  if (compressed_data == NULL) {
+    log_message(LOG_LEVEL_ERROR, "Failed to allocate compressed data");
+    goto cleanup;
+  }
+  if (output.pos > 0)
+    memcpy(compressed_data->data, ctx->out_buf, output.pos);
   compressed_data->size = output.pos;
-  ZSTD_freeCCtx(cctx);
 
   log_debug_message(LOG_DEBUG_UTIL, "Data succesfully compressed from %zu to %zu",
                     data_to_compress->size, compressed_data->size);
+
+cleanup:
+  compression_ctx_put(ctx);
   return compressed_data;
 }
 
@@ -144,20 +269,30 @@ Data* data_decompress_limited(Data* compressed_data, size_t maximum_size) {
     return NULL;
   }
 
-  ZSTD_DCtx* dctx = ZSTD_createDCtx();
-  if (!dctx) {
-    log_message(LOG_LEVEL_ERROR, "Failed to create ZSTD decompression context");
+  CompressionThreadCtx* ctx = compression_get_thread_ctx();
+  if (ctx == NULL) {
+    log_message(LOG_LEVEL_ERROR, "Failed to allocate ZSTD decompression context");
     return NULL;
   }
+  Data* uncompressed_data = NULL;
+
+  if (!ctx->dctx) {
+    ctx->dctx = ZSTD_createDCtx();
+    if (!ctx->dctx) {
+      log_message(LOG_LEVEL_ERROR, "Failed to create ZSTD decompression context");
+      goto cleanup;
+    }
+  }
+  /* Reset only the session; decompression parameters are sticky. */
+  ZSTD_DCtx_reset(ctx->dctx, ZSTD_reset_session_only);
 
   size_t buf_size = (dst_size > 0) ? (size_t)dst_size : INITIAL_DECOMPRESS_BUF_SIZE;
   if (buf_size > maximum_size)
     buf_size = maximum_size;
-  Data* uncompressed_data = data_create_empty(buf_size);
+  uncompressed_data = data_create_empty(buf_size);
   if (!uncompressed_data) {
     log_message(LOG_LEVEL_ERROR, "Failed to allocate decompression buffer");
-    ZSTD_freeDCtx(dctx);
-    return NULL;
+    goto cleanup;
   }
 
   ZSTD_inBuffer input = {compressed_data->data, compressed_data->size, 0};
@@ -165,20 +300,20 @@ Data* data_decompress_limited(Data* compressed_data, size_t maximum_size) {
 
   size_t ret;
   do {
-    ret = ZSTD_decompressStream(dctx, &output, &input);
+    ret = ZSTD_decompressStream(ctx->dctx, &output, &input);
     if (ZSTD_isError(ret)) {
       log_message(LOG_LEVEL_ERROR, "Decompression failed: %s", ZSTD_getErrorName(ret));
-      ZSTD_freeDCtx(dctx);
       data_destroy(uncompressed_data);
-      return NULL;
+      uncompressed_data = NULL;
+      goto cleanup;
     }
     if (ret > 0 && output.pos == output.size) {
       if (buf_size >= hard_limit || buf_size > SIZE_MAX / 2) {
         log_message(LOG_LEVEL_ERROR, "Decompressed data exceeds %llu bytes",
                     (unsigned long long)MAX_DECOMPRESSED_SIZE);
-        ZSTD_freeDCtx(dctx);
         data_destroy(uncompressed_data);
-        return NULL;
+        uncompressed_data = NULL;
+        goto cleanup;
       }
       buf_size *= 2;
       if (buf_size > hard_limit)
@@ -186,9 +321,9 @@ Data* data_decompress_limited(Data* compressed_data, size_t maximum_size) {
       void* new_data = protocol_realloc(uncompressed_data->data, buf_size);
       if (!new_data) {
         log_message(LOG_LEVEL_ERROR, "Failed to grow decompression buffer");
-        ZSTD_freeDCtx(dctx);
         data_destroy(uncompressed_data);
-        return NULL;
+        uncompressed_data = NULL;
+        goto cleanup;
       }
       uncompressed_data->data = new_data;
       output.dst = new_data;
@@ -197,9 +332,11 @@ Data* data_decompress_limited(Data* compressed_data, size_t maximum_size) {
   } while (ret > 0);
 
   uncompressed_data->size = output.pos;
-  ZSTD_freeDCtx(dctx);
 
   log_debug_message(LOG_DEBUG_UTIL, "Decompressed data successfully");
+
+cleanup:
+  compression_ctx_put(ctx);
   return uncompressed_data;
 }