src/client/client_send.c

@@ -1157,6 +1157,7 @@ static int send_append(const Client* client, File* file, Config* config,
     tail_view.data = (char*)file->data->data + off;
     tail_view.size = tail_len;
     tail_view.protocol_charge = 0;
+    tail_view.owner = NULL;
     ok = send_data(fd, &tail_view);
   }
   return ok ? 0 : -1;

src/shared/data.c

@@ -23,6 +23,7 @@ Data* data_create_reserve(size_t size) {
   d->data = NULL;
   d->size = size;
   d->protocol_charge = 0;
+  d->owner = NULL;
   return d;
 }
 
@@ -36,14 +37,19 @@ Data* data_create(void* data, size_t data_size) {
   new_data->data = data;
   new_data->size = data_size;
   new_data->protocol_charge = 0;
+  new_data->owner = NULL;
   return new_data;
 }
 
 void data_destroy(Data* data) {
   if (data == NULL)
     return;
-  if (data->protocol_charge != 0)
+  if (data->protocol_charge != 0) {
+    if (data->owner != NULL)
+      protocol_release_memory_for_session(data->owner, data->protocol_charge);
+    else
       protocol_release_memory(data->protocol_charge);
+  }
   free(data->data);
   free(data);
 }

src/shared/data.h

@@ -3,11 +3,19 @@
 
 #include <stdlib.h>
 
+/* Forward declaration for the connection budget a received Data is charged
+ * against; defined in protocol.h (which includes this header). */
+typedef struct ProtocolSession ProtocolSession;
+
 typedef struct {
   void* data;
   size_t size;
   /* Non-zero only for a buffer charged to the protocol connection budget. */
   size_t protocol_charge;
+  /* Session whose budget `protocol_charge` was reserved from.  The charge must
+   * always be returned to this session, regardless of which session (if any) is
+   * bound to the destroying thread.  NULL for uncharged Data. */
+  ProtocolSession* owner;
 } Data;
 
 Data* data_create_empty(size_t data_size);
@@ -15,5 +23,8 @@ Data* data_create_reserve(size_t size);
 Data* data_create(void* data, size_t data_size);
 void data_destroy(Data* data);
 void protocol_release_memory(size_t charge);
+/* Release `charge` against `session` directly instead of the thread-local bound
+ * session.  Used by data_destroy to honor Data.owner. */
+void protocol_release_memory_for_session(ProtocolSession* session, size_t charge);
 
 #endif

src/shared/protocol.c

@@ -40,7 +40,7 @@ static bool protocol_reserve_memory(ProtocolSession* session, size_t charge) {
   }
 }
 
-static void protocol_release_memory_for_session(ProtocolSession* session, size_t charge) {
+void protocol_release_memory_for_session(ProtocolSession* session, size_t charge) {
   unsigned long long allocated = atomic_load(&session->total_allocated_bytes);
   while (true) {
     unsigned long long remaining = (unsigned long long)charge >= allocated ? 0 : allocated - charge;
@@ -573,6 +573,7 @@ Data* protocol_receive_data_limited(ProtocolSession* session, unsigned long long
     return NULL;
   }
   result->protocol_charge = allocation_size;
+  result->owner = session;
   return result;
 }
 

tests/test_protocol.c

@@ -412,6 +412,54 @@ static void test_protocol_accounting_release_does_not_underflow() {
   protocol_session_unbind();
 }
 
+/* A Data acquired on session A must return its connection-memory charge to A
+   even when a different session B is bound at destroy time: releasing against
+   the thread-local bound session would leak A's budget and drain B's. */
+static void test_receive_data_charge_follows_owning_session() {
+  int pipe_a[2];
+  int pipe_b[2];
+  EXPECT_EQ_INT(pipe(pipe_a), 0);
+  EXPECT_EQ_INT(pipe(pipe_b), 0);
+
+  ProtocolSession session_a;
+  ProtocolSession session_b;
+  protocol_session_init(&session_a, pipe_a[0], pipe_a[1]);
+  protocol_session_init(&session_b, pipe_b[0], pipe_b[1]);
+  protocol_session_set_max_alloc(&session_a, 64);
+  protocol_session_set_max_alloc(&session_b, 64);
+
+  unsigned long long size = 8;
+  EXPECT_EQ_INT((int)write(pipe_a[1], &size, sizeof(size)), (int)sizeof(size));
+  EXPECT_EQ_INT((int)write(pipe_a[1], "12345678", 8), 8);
+  EXPECT_EQ_INT((int)write(pipe_b[1], &size, sizeof(size)), (int)sizeof(size));
+  EXPECT_EQ_INT((int)write(pipe_b[1], "abcdefgh", 8), 8);
+
+  Data* data_a = protocol_receive_data_limited(&session_a, 8);
+  Data* data_b = protocol_receive_data_limited(&session_b, 8);
+  EXPECT_NOT_NULL(data_a);
+  EXPECT_NOT_NULL(data_b);
+  EXPECT_TRUE(data_a->owner == &session_a);
+  EXPECT_TRUE(data_b->owner == &session_b);
+  EXPECT_EQ_INT((int)atomic_load(&session_a.total_allocated_bytes), 8);
+  EXPECT_EQ_INT((int)atomic_load(&session_b.total_allocated_bytes), 8);
+
+  /* Destroy A's Data while the unrelated session B is the bound session. */
+  protocol_session_bind(&session_b);
+  data_destroy(data_a);
+  protocol_session_unbind();
+
+  EXPECT_EQ_INT((int)atomic_load(&session_a.total_allocated_bytes), 0);
+  EXPECT_EQ_INT((int)atomic_load(&session_b.total_allocated_bytes), 8);
+
+  data_destroy(data_b);
+  EXPECT_EQ_INT((int)atomic_load(&session_b.total_allocated_bytes), 0);
+
+  close(pipe_a[0]);
+  close(pipe_a[1]);
+  close(pipe_b[0]);
+  close(pipe_b[1]);
+}
+
 static void test_protocol_session_io_timeout() {
   /* Default is the built-in 60 s window; the setter stores exactly what it is
    * given (<= 0 means "fall back to the default") so callers can propagate
@@ -574,4 +622,5 @@ void test_protocol() {
   test_protocol_accounting_reservation_is_atomic();
   test_protocol_string_accounting_is_transient();
   test_protocol_accounting_release_does_not_underflow();
+  test_receive_data_charge_follows_owning_session();
 }