Merge branch 'fix/w4-queue' into fix/w4-integration

src/shared/multiprocessing.c

@@ -11,6 +11,7 @@
 #include "protocol.h"
 #include "queue.h"
 #include "utils.h"
+#include <stdint.h>
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
@@ -26,6 +27,8 @@ PipelineContextSender* pipeline_context_sender_create(Config* config, Queue* que
   context->queue_loader = queue_loader;
   context->scanner_done = false;
   context->loader_done = false;
+  context->queued_bytes = 0;
+  context->max_queue_bytes = 0;
   context->manifest = NULL;
   context->excluded_paths = NULL;
   context->missing_args = NULL;
@@ -97,6 +100,84 @@ fail:
   return NULL;
 }
 
+void pipeline_context_sender_set_queue_byte_limit(PipelineContextSender* context,
+                                                  size_t max_bytes) {
+  if (context == NULL)
+    return;
+  mtx_lock(&context->mutex_loader);
+  context->max_queue_bytes = max_bytes;
+  context->queued_bytes = 0;
+  cnd_broadcast(&context->condition_not_full_loader);
+  mtx_unlock(&context->mutex_loader);
+}
+
+size_t pipeline_context_sender_chunk_bytes(const Chunk* chunk) {
+  if (chunk == NULL || chunk->items == NULL)
+    return 0;
+  size_t total = 0;
+  for (int i = 0; i < chunk->element_count; i++) {
+    const File* file = chunk->items[i];
+    if (file == NULL || file->data == NULL || file->data->data == NULL)
+      continue;
+    if (file->data->size > SIZE_MAX - total)
+      return SIZE_MAX;
+    total += file->data->size;
+  }
+  return total;
+}
+
+void pipeline_context_sender_note_bytes_released(PipelineContextSender* context,
+                                                 size_t released_bytes) {
+  if (context == NULL || context->max_queue_bytes == 0 || released_bytes == 0)
+    return;
+  mtx_lock(&context->mutex_loader);
+  if (released_bytes >= context->queued_bytes)
+    context->queued_bytes = 0;
+  else
+    context->queued_bytes -= released_bytes;
+  cnd_signal(&context->condition_not_full_loader);
+  mtx_unlock(&context->mutex_loader);
+}
+
+bool pipeline_context_sender_enqueue_chunk(PipelineContextSender* context, Chunk* chunk) {
+  if (context == NULL || chunk == NULL)
+    return false;
+  size_t chunk_bytes = pipeline_context_sender_chunk_bytes(chunk);
+  mtx_lock(&context->mutex_loader);
+  while (!atomic_load(&context->cancelled)) {
+    bool blocked_by_count = queue_is_full(context->queue_loader);
+    bool blocked_by_budget = false;
+    if (context->max_queue_bytes > 0) {
+      size_t budget = context->max_queue_bytes;
+      size_t used = context->queued_bytes;
+      if (used >= budget) {
+        blocked_by_budget = true;
+      } else if (chunk_bytes > budget - used) {
+        /* A single payload larger than the whole budget is only admitted to an
+           empty pipeline so the wait can never deadlock. */
+        blocked_by_budget = used != 0;
+      }
+    }
+    if (!blocked_by_count && !blocked_by_budget)
+      break;
+    cnd_wait(&context->condition_not_full_loader, &context->mutex_loader);
+  }
+  if (atomic_load(&context->cancelled)) {
+    mtx_unlock(&context->mutex_loader);
+    chunk_destroy(chunk);
+    return false;
+  }
+  if (!queue_enqueue(context->queue_loader, chunk)) {
+    mtx_unlock(&context->mutex_loader);
+    chunk_destroy(chunk);
+    return false;
+  }
+  context->queued_bytes += chunk_bytes;
+  cnd_signal(&context->condition_not_empty_loader);
+  mtx_unlock(&context->mutex_loader);
+  return true;
+}
+
 void pipeline_context_sender_destroy(PipelineContextSender* context) {
   if (context->manifest) {
     array_list_delete(context->manifest);

src/shared/multiprocessing.h

@@ -5,6 +5,7 @@
 #include <stdatomic.h>
 
 #include "array_list.h"
+#include "chunk.h"
 #include "config.h"
 #include "file.h"
 #include "protocol.h"
@@ -25,6 +26,15 @@ typedef struct {
   cnd_t condition_not_full_loader;
   cnd_t condition_not_empty_loader;
   bool loader_done;
+  /* Aggregate loaded payload bytes queued on queue_loader but not yet released
+     by the sender.  Guarded by `mutex_loader`.  When `max_queue_bytes` is
+     non-zero the loader blocks before enqueueing a chunk that would push this
+     total over it, so the sender buffers a bounded number of bytes rather than
+     an unbounded count of chunks that may each be up to chunk_size (or a single
+     file) in size.  Files streamed straight from disk by sendfile hold no
+     payload, so only in-memory (`data->data`) payloads are counted. */
+  size_t queued_bytes;
+  size_t max_queue_bytes;
   ArrayList* manifest;
   /* Protected prefixes (paths the source scan excluded by user rules) sent
      with the keep-set manifest so --delete leaves them alone unless
@@ -113,6 +123,22 @@ typedef struct PipelineContextReceiver {
 PipelineContextSender* pipeline_context_sender_create(Config* config, Queue* queue_scanner,
                                                       Queue* queue_loader);
 void pipeline_context_sender_destroy(PipelineContextSender* context);
+/* Bound the loaded payload bytes the sender may buffer ahead of the network
+   writer (see max_queue_bytes). */
+void pipeline_context_sender_set_queue_byte_limit(PipelineContextSender* context, size_t max_bytes);
+/* Total payload bytes a chunk currently holds in memory (loaded file data
+   only; zero for entries with no payload or data streamed from disk). */
+size_t pipeline_context_sender_chunk_bytes(const Chunk* chunk);
+/* Blocking enqueue used by the sender's loader stage.  Blocks while
+   queue_loader is full by element count or when adding `chunk` would push the
+   queued payload bytes over the configured byte limit; waits until the sender
+   releases bytes.  Takes ownership of `chunk` on success and destroys it on
+   failure/cancel. */
+bool pipeline_context_sender_enqueue_chunk(PipelineContextSender* context, Chunk* chunk);
+/* Account for `released_bytes` of payload memory that the sender freed after
+   destroying a chunk, unblocking a loader waiting on the byte limit. */
+void pipeline_context_sender_note_bytes_released(PipelineContextSender* context,
+                                                 size_t released_bytes);
 PipelineContextReceiver* pipeline_context_receiver_create(Config* config, Queue* queue_receiver,
                                                           int file_descriptor, SSL* ssl);
 void pipeline_context_receiver_destroy(PipelineContextReceiver* context);