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;
 }
 

src/shared/compression.h

@@ -14,4 +14,12 @@ Data* data_decompress_limited(Data* compressed_data, size_t maximum_size);
 bool compression_should_skip(const char* path);
 bool compression_should_skip_with_suffixes(const char* path, char* const* suffixes, int count);
 
+/* Release the calling thread's cached zstd contexts (compressor, decompressor
+ * and scratch buffer).  The cache is thread-local and is also released
+ * automatically when a worker thread exits (via a C11 tss destructor) and for
+ * the main thread at process exit; this explicit entry point exists so tests
+ * and long-lived callers can drop the cache deterministically.  Safe to call
+ * when no context has been created, and idempotent. */
+void compression_free_thread_contexts(void);
+
 #endif

src/shared/transport_tcp.c

@@ -19,6 +19,7 @@
 static volatile sig_atomic_t g_active_connections = 0;
 
 static void tcp_apply_socket_timeout(int fd);
+static void tcp_enable_nodelay_default(int fd, int family);
 
 static void sigchld_handler(int sig) {
   (void)sig;
@@ -148,6 +149,7 @@ static void accept_loop(Server* server, void (*child_fn)(int, void*), void* chil
       continue;
     }
     tcp_apply_socket_timeout(fd);
+    tcp_enable_nodelay_default(fd, client_addr.ss_family);
     char peer[128];
     if (!utils_sockaddr_to_string((const struct sockaddr*)&client_addr, peer, sizeof(peer)))
       snprintf(peer, sizeof(peer), "unknown");
@@ -231,6 +233,19 @@ static void tcp_apply_socket_timeout(int fd) {
   setsockopt(fd, SOL_SOCKET, SO_SNDTIMEO, &tv, sizeof(tv));
 }
 
+/* Enable TCP_NODELAY by default on a transfer socket: the protocol emits many
+ * small messages and Nagle's algorithm would otherwise coalesce/delay them.
+ * Best-effort only: the family guard keeps this to IP/TCP sockets, and a
+ * setsockopt failure is ignored.  A caller-provided --sockopts TCP_NODELAY=0
+ * is applied afterwards on the connect path, so an explicit user choice still
+ * wins. */
+static void tcp_enable_nodelay_default(int fd, int family) {
+  if (family != AF_INET && family != AF_INET6)
+    return;
+  int value = 1;
+  setsockopt(fd, IPPROTO_TCP, TCP_NODELAY, &value, sizeof(value));
+}
+
 Client* client_create() {
   Client* client = (Client*)malloc(sizeof(Client));
   if (client == NULL) {
@@ -376,6 +391,9 @@ bool tcp_connect_socket_ex(Client* client, const char* host, int port,
     if (client->file_descriptor < 0)
       continue;
 
+    /* Default first; a user --sockopts TCP_NODELAY=0 applied below overrides. */
+    tcp_enable_nodelay_default(client->file_descriptor, rp->ai_family);
+
     if (opts && opts->sockopt_count > 0 &&
         !tcp_apply_sockopts(client->file_descriptor, opts->sockopts, opts->sockopt_count)) {
       close(client->file_descriptor);

tests/test_compression.c

@@ -6,6 +6,7 @@
 #include "utils.h"
 #include <string.h>
 #include <sys/stat.h>
+#include <threads.h>
 #include <unistd.h>
 
 static void test_data_compress_decompress_roundtrip() {
@@ -137,10 +138,84 @@ static void test_chunk_compress_decompress_roundtrip() {
   unlink(path2);
 }
 
+typedef struct {
+  int id;
+  int iterations;
+  bool ok;
+} CompressionThreadArg;
+
+/* Each worker exercises the per-thread cached zstd contexts: several
+ * compress/decompress round-trips with varying payload sizes, levels and
+ * worker counts so the context is reused (and its parameters re-applied)
+ * across calls, concurrently with other workers. */
+static int compression_reuse_worker(void* arg) {
+  CompressionThreadArg* a = (CompressionThreadArg*)arg;
+  a->ok = true;
+  for (int it = 0; it < a->iterations; it++) {
+    size_t size = 512 + (size_t)((a->id * 7919 + it * 104729) % (48 * 1024));
+    char* original = malloc(size);
+    if (!original) {
+      a->ok = false;
+      break;
+    }
+    for (size_t i = 0; i < size; i++)
+      original[i] = (char)((i * 31 + (size_t)a->id + (size_t)it * 7) % 251);
+    Data* input = data_create(original, size);
+    if (!input) { /* data_create takes ownership of original, even on failure */
+      a->ok = false;
+      break;
+    }
+    int level = 1 + ((it / 2) % 5);
+    int threads = ((it / 2) % 2 == 0) ? 2 : 0;
+    Data* compressed = data_compress_with_threads(input, level, threads);
+    if (!compressed) {
+      data_destroy(input);
+      a->ok = false;
+      break;
+    }
+    Data* decompressed = data_decompress(compressed);
+    bool roundtrip_ok = decompressed != NULL && decompressed->size == size &&
+                        memcmp(decompressed->data, original, size) == 0;
+    data_destroy(decompressed);
+    data_destroy(compressed);
+    data_destroy(input);
+    if (!roundtrip_ok) {
+      a->ok = false;
+      break;
+    }
+  }
+  /* Deliberately do NOT free the thread context here: the C11 tss destructor
+   * must release it when this thread exits (validated by LeakSanitizer). */
+  return thrd_success;
+}
+
+static void test_data_compress_reused_contexts_multithreaded() {
+  enum { NTHREADS = 8, ITERATIONS = 6 };
+  thrd_t threads[NTHREADS];
+  CompressionThreadArg args[NTHREADS];
+  bool all_created = true;
+  for (int i = 0; i < NTHREADS; i++) {
+    args[i].id = i;
+    args[i].iterations = ITERATIONS;
+    args[i].ok = false;
+    if (thrd_create(&threads[i], compression_reuse_worker, &args[i]) != thrd_success) {
+      all_created = false;
+      break;
+    }
+  }
+  EXPECT_TRUE(all_created);
+  for (int i = 0; i < NTHREADS; i++)
+    EXPECT_EQ_INT(thrd_join(threads[i], NULL), thrd_success);
+  for (int i = 0; i < NTHREADS; i++)
+    EXPECT_TRUE(args[i].ok);
+  compression_free_thread_contexts();
+}
+
 void test_compression() {
   test_data_compress_decompress_roundtrip();
   test_data_compress_decompress_large();
   test_skip_compress_suffix_matching();
   test_data_compress_with_threads_roundtrip();
+  test_data_compress_reused_contexts_multithreaded();
   test_chunk_compress_decompress_roundtrip();
 }

tests/test_transport_tcp.c

@@ -2,6 +2,7 @@
 #include "protocol.h"
 #include "test_utils.h"
 #include "transport_tcp.h"
+#include <arpa/inet.h>
 #include <netinet/in.h>
 #include <netinet/tcp.h>
 #include <string.h>
@@ -195,6 +196,49 @@ static void test_client_disconnect_delete() {
   client_delete(c);
 }
 
+/* TCP_NODELAY is enabled by default on a connected transfer socket, and an
+ * explicit --sockopts TCP_NODELAY=0 still overrides it. */
+static void test_tcp_nodelay_default_and_override() {
+  Server* s = server_create(0);
+  EXPECT_NOT_NULL(s);
+  EXPECT_EQ_INT(listen(s->file_descriptor, 1), 0);
+  struct sockaddr_in bound;
+  socklen_t bound_len = sizeof(bound);
+  EXPECT_EQ_INT(getsockname(s->file_descriptor, (struct sockaddr*)&bound, &bound_len), 0);
+  int port = (int)ntohs(bound.sin_port);
+  EXPECT_TRUE(port > 0);
+
+  Client* c = client_create();
+  EXPECT_NOT_NULL(c);
+  EXPECT_TRUE(client_connect(c, "127.0.0.1", port));
+  int got = 0;
+  socklen_t len = sizeof(got);
+  EXPECT_EQ_INT(getsockopt(c->file_descriptor, IPPROTO_TCP, TCP_NODELAY, &got, &len), 0);
+  EXPECT_EQ_INT(got, 1);
+  client_disconnect(c);
+  client_delete(c);
+
+  SockOptEntry* entries = NULL;
+  int count = 0;
+  EXPECT_EQ_INT(config_sockopts_parse("TCP_NODELAY=0", &entries, &count), 0);
+  TcpConnectOptions opts;
+  memset(&opts, 0, sizeof(opts));
+  opts.sockopts = entries;
+  opts.sockopt_count = count;
+
+  Client* c2 = client_create();
+  EXPECT_NOT_NULL(c2);
+  EXPECT_TRUE(client_connect_ex(c2, "127.0.0.1", port, &opts));
+  got = 0;
+  len = sizeof(got);
+  EXPECT_EQ_INT(getsockopt(c2->file_descriptor, IPPROTO_TCP, TCP_NODELAY, &got, &len), 0);
+  EXPECT_EQ_INT(got, 0);
+  client_disconnect(c2);
+  client_delete(c2);
+  free(entries);
+  server_delete(&s);
+}
+
 void test_transport_tcp() {
   test_server_create_ephemeral();
   test_server_delete_null();
@@ -211,4 +255,5 @@ void test_transport_tcp() {
   test_sockopts_apply_sets_option();
   test_server_create_bind_address();
   test_server_create_bind_ipv6();
+  test_tcp_nodelay_default_and_override();
 }