fix(mem): hand workers their freed blocks back at the end of every dispatch

src/core/pool.c

@@ -82,6 +82,10 @@ static void worker_loop(void* arg) {
      * by an earlier worker (pools and slab caches intact). */
     ray_heap_init();
     ray_rc_sync = true;  /* workers always use atomic refcounting */
+    /* Publish the heap for the dispatcher's post-dispatch reclaim (release
+     * pairs with its acquire: the heap is fully initialised when seen). */
+    atomic_store_explicit(&pool->worker_heaps[wctx.worker_id - 1], ray_tl_heap,
+                          memory_order_release);
 
     for (;;) {
         ray_sem_wait(&pool->work_ready);
@@ -109,9 +113,10 @@ static void worker_loop(void* arg) {
                                       memory_order_acq_rel);
         }
 
-        /* No ray_heap_gc() here — removing worker GC between dispatch rounds
-         * ensures main can safely modify worker heaps in ray_parallel_end().
-         * Eager madvise in heap_coalesce already releases pages on free. */
+        /* No ray_heap_gc() here — a worker that neither allocates nor frees
+         * between its last pending-- and sem_wait is what lets the
+         * dispatcher drain worker foreign lists (ray_heap_reclaim_workers)
+         * and the idle decay walk worker heaps once the flag is clear. */
     }
 
     /* Abandon, do not destroy.  Another thread may still hold — and later
@@ -120,9 +125,26 @@ static void worker_loop(void* arg) {
      * a heap that could be unregistered and munmapped underneath that lookup
      * would be a use-after-free.  The heap stays registered with its pools
      * and is adopted by the next worker thread that starts. */
+    atomic_store_explicit(&pool->worker_heaps[wctx.worker_id - 1], NULL,
+                          memory_order_release);
     ray_heap_abandon();
 }
 
+/* End of a parallel region: hand every worker the blocks other threads
+ * freed to it since the last dispatch, so the next round reuses them
+ * instead of cutting fresh pool space (issue #619).  Workers only — they
+ * are parked on the semaphore and touch nothing of their own until the
+ * next dispatch; any other registered heap may belong to a live thread.
+ * The dispatcher's own list is drained too, on its own thread. */
+static void pool_reclaim_worker_heaps(ray_pool_t* pool) {
+    for (uint32_t i = 0; i < pool->n_workers; i++) {
+        ray_heap_t* h = (ray_heap_t*)atomic_load_explicit(&pool->worker_heaps[i],
+                                                          memory_order_acquire);
+        if (h) ray_heap_reclaim_worker(h);
+    }
+    ray_heap_flush_foreign();
+}
+
 /* --------------------------------------------------------------------------
  * ray_pool_create
  * -------------------------------------------------------------------------- */
@@ -192,6 +214,15 @@ static ray_err_t ray_pool_create_impl(ray_pool_t* pool, uint32_t n_workers,
             ray_sys_free(pool->tasks);
             return RAY_ERR_OOM;
         }
+        pool->worker_heaps = ray_sys_alloc(n_workers * sizeof(*pool->worker_heaps));
+        if (!pool->worker_heaps) {
+            ray_sys_free(pool->threads);
+            ray_sem_destroy(&pool->work_ready);
+            ray_sys_free(pool->tasks);
+            return RAY_ERR_OOM;
+        }
+        for (uint32_t i = 0; i < n_workers; i++)
+            atomic_store_explicit(&pool->worker_heaps[i], NULL, memory_order_relaxed);
 
         for (uint32_t i = 0; i < n_workers; i++) {
             worker_ctx_t* wctx = (worker_ctx_t*)ray_sys_alloc(sizeof(worker_ctx_t));
@@ -204,6 +235,7 @@ static ray_err_t ray_pool_create_impl(ray_pool_t* pool, uint32_t n_workers,
                 for (uint32_t j = 0; j < i; j++) {
                     ray_thread_join(pool->threads[j]);
                 }
+                ray_sys_free(pool->worker_heaps);
                 ray_sys_free(pool->threads);
                 ray_sem_destroy(&pool->work_ready);
                 ray_sys_free(pool->tasks);
@@ -222,6 +254,7 @@ static ray_err_t ray_pool_create_impl(ray_pool_t* pool, uint32_t n_workers,
                 for (uint32_t j = 0; j < i; j++) {
                     ray_thread_join(pool->threads[j]);
                 }
+                ray_sys_free(pool->worker_heaps);
                 ray_sys_free(pool->threads);
                 ray_sem_destroy(&pool->work_ready);
                 ray_sys_free(pool->tasks);
@@ -255,6 +288,7 @@ void ray_pool_free(ray_pool_t* pool) {
         ray_thread_join(pool->threads[i]);
     }
 
+    ray_sys_free(pool->worker_heaps);
     ray_sys_free(pool->threads);
     ray_sem_destroy(&pool->work_ready);
     ray_sys_free(pool->tasks);
@@ -392,6 +426,8 @@ void ray_pool_dispatch(ray_pool_t* pool, ray_pool_fn fn, void* ctx,
      * be between pending-- and sem_wait. */
     atomic_thread_fence(memory_order_seq_cst);
     ray_rc_sync = false;
+
+    pool_reclaim_worker_heaps(pool);
 }
 
 /* One round of ray_pool_dispatch_n: tasks [first, first+n_tasks), each handed
@@ -466,6 +502,7 @@ static void dispatch_n_round(ray_pool_t* pool, ray_pool_fn fn, void* ctx,
     atomic_store_explicit(&ray_parallel_flag, 0, memory_order_release);
     atomic_thread_fence(memory_order_seq_cst);
     ray_rc_sync = false;
+    pool_reclaim_worker_heaps(pool);
 }
 
 /* --------------------------------------------------------------------------

src/core/pool.h

@@ -52,6 +52,16 @@ struct ray_pool {
     uint32_t           n_workers;     /* number of background threads (nproc - 1) */
     _Atomic(uint32_t)  shutdown;
 
+    /* Heap of each worker thread [n_workers] (a ray_heap_t*), published by
+     * the worker once ray_heap_init has run and cleared before it exits.
+     * The dispatcher reads them at the end of every parallel region to hand
+     * each worker the blocks other threads freed to it (ray_heap_reclaim_
+     * worker).  Only these heaps: a worker parked on the semaphore is the
+     * one thread known to touch nothing of its own until the next dispatch,
+     * which is what makes draining its list from here sound.  NULL until
+     * the worker is up. */
+    _Atomic(void*)*     worker_heaps;
+
     /* SPMC task ring (single producer = main, multi consumer = workers + main).
      *
      * Claiming uses two MONOTONIC 64-bit cursors that are never reset, which

src/mem/heap.c

@@ -2793,6 +2793,44 @@ void ray_heap_flush_foreign(void) {
     heap_drain_foreign(h);
 }
 
+/* --------------------------------------------------------------------------
+ * Post-dispatch reclaim (dispatcher side)
+ *
+ * A parallel operator's workers allocate their per-task buffers from their
+ * own heaps and the main thread frees them once the dispatch has completed,
+ * so every such block ends up on the owning worker's foreign list.  The
+ * owner takes that list back only when an allocation finds its freelists
+ * empty at the order it needs — and a warm worker with a partly cut pool
+ * rarely does: it keeps splitting fresh pool space instead, touching new
+ * pages every round while its own freed blocks wait on the list.  Under a
+ * steady stream of such operators the process grows by one full pool per
+ * worker before any block is reused, and nothing short of the idle decay
+ * (which needs the process to sit quiet) drains it earlier.
+ *
+ * So the dispatcher drains each worker's list at the end of each parallel
+ * region.  The conditions are the ones the decay sweep relies on: the flag
+ * is clear, so every worker has done its last pending-- and is claiming
+ * nothing, allocating nothing and freeing nothing until the next dispatch;
+ * a concurrent push onto a foreign list from some other thread is safe
+ * because the drain takes the whole list in one exchange and leaves later
+ * arrivals for the next round.  Only pool workers qualify: the registry
+ * also holds the heaps of other live threads (a server's poll thread, an
+ * embedding's own threads), whose freelists only their owner may touch, so
+ * the pool names the heaps rather than this walking the registry.  No
+ * pages are released here — the blocks only go back to freelists, so the
+ * next round reuses them without faulting.  Cost: one load per worker and
+ * the coalescing of whatever was freed cross-thread since the last
+ * dispatch, work the owner would otherwise do on its next dry allocation.
+ * -------------------------------------------------------------------------- */
+
+void ray_heap_reclaim_worker(ray_heap_t* h) {
+    if (!h) return;
+    if (atomic_load_explicit(&ray_parallel_flag, memory_order_acquire) != 0)
+        return;
+    if (!atomic_load_explicit(&h->foreign, memory_order_relaxed)) return;
+    heap_drain_foreign(h);
+}
+
 /* --------------------------------------------------------------------------
  * Pending-merge queue (lock-free LIFO)
  *

src/mem/heap.h

@@ -226,6 +226,13 @@ void     ray_heap_destroy(void);
 void     ray_heap_abandon(void);
 void     ray_heap_merge(ray_heap_t* src);
 void     ray_heap_flush_foreign(void);
+/* Give a parked worker's heap the blocks other threads freed to it.  For the
+ * dispatcher, once a parallel region has ended: a worker only drains its own
+ * list when its freelists run dry, so between dispatches the blocks the main
+ * thread freed sit unused while the worker keeps cutting fresh pool space.
+ * The caller vouches that the owning thread is idle (a pool worker on its
+ * semaphore); must run with ray_parallel_flag clear, does nothing otherwise. */
+void     ray_heap_reclaim_worker(ray_heap_t* h);
 void     ray_heap_push_pending(ray_heap_t* heap);
 void     ray_heap_drain_pending(void);
 uint8_t  ray_order_for_size(size_t data_size);

test/test_heap.c

@@ -41,6 +41,7 @@
 #include "test.h"
 #include <rayforce.h>
 #include "mem/heap.h"
+#include "core/platform.h"
 #include "core/pool.h"
 
 /* Restore the shipped policy after a test drives it. */
@@ -640,6 +641,67 @@ static test_result_t test_free_routes_to_owner_list(void) {
     PASS();
 }
 
+/* ---- The dispatcher hands a foreign block back to its owner --------------
+ *
+ * Issue #619: a parallel operator's workers allocate per-task buffers and the
+ * main thread frees them after the dispatch, so they land on the owning
+ * worker's foreign list, where they sit until that worker's freelists run
+ * dry — which a warm worker's rarely do.  ray_heap_reclaim_worker, run by
+ * the dispatcher on each worker's heap at the end of every parallel region,
+ * drains that list into the worker's freelists.  Model it with two heaps: a
+ * block owned by heap_b and freed from heap_a stays parked (still on heap_b's
+ * books) until heap_a reclaims heap_b — and is left alone while the parallel
+ * flag is set. */
+
+static test_result_t test_reclaim_worker_drains_owner_list(void) {
+    ray_heap_t* heap_a = ray_tl_heap;
+
+    ray_tl_heap = NULL;
+    ray_heap_init();
+    ray_heap_t* heap_b = ray_tl_heap;
+    TEST_ASSERT_NOT_NULL(heap_b);
+    /* heap_b may be a heap abandoned by an earlier test's worker, with blocks
+     * freed to it since; take them back first so the books below move only
+     * for the block this test frees. */
+    ray_heap_flush_foreign();
+#if RAY_MEM_STATS
+    size_t booked0 = heap_b->stats.bytes_allocated;
+#endif
+    ray_t* blk = ray_alloc(256u << 10);          /* above the slab orders */
+    TEST_ASSERT_NOT_NULL(blk);
+#if RAY_MEM_STATS
+    size_t booked1 = heap_b->stats.bytes_allocated;
+    TEST_ASSERT(booked1 > booked0, "allocation charged to heap_b");
+#endif
+
+    ray_tl_heap = heap_a;
+    ray_free(blk);
+    TEST_ASSERT_EQ_U((uintptr_t)atomic_load(&heap_b->foreign), (uintptr_t)blk);
+#if RAY_MEM_STATS
+    /* Parked, not yet returned: heap_b still carries the bytes. */
+    TEST_ASSERT_EQ_U(heap_b->stats.bytes_allocated, booked1);
+#endif
+
+    /* Inside a parallel region the reclaim is a no-op. */
+    ray_parallel_begin();
+    ray_heap_reclaim_worker(heap_b);
+    TEST_ASSERT_EQ_U((uintptr_t)atomic_load(&heap_b->foreign), (uintptr_t)blk);
+    atomic_store(&ray_parallel_flag, 0);
+
+    /* Once it is over, another thread reclaims heap_b: the list is empty and
+     * the block is back on heap_b's freelists, off its books. */
+    ray_heap_reclaim_worker(heap_b);
+    TEST_ASSERT_NULL(atomic_load(&heap_b->foreign));
+#if RAY_MEM_STATS
+    TEST_ASSERT_EQ_U(heap_b->stats.bytes_allocated, booked0);
+#endif
+
+    ray_tl_heap = heap_b;
+    ray_heap_destroy();
+    ray_tl_heap = heap_a;
+    PASS();
+}
+
 /* ---- A cross-thread free cycle must not grow the owner's pools ----------
  *
  * The regression for issue #439.  Workers allocate per-group temporaries from
@@ -2961,6 +3023,7 @@ const test_entry_t heap_entries[] = {
     { "heap/gc_serial",                test_heap_gc_serial,              heap_setup, heap_teardown },
     { "heap/gc_parallel",              test_heap_gc_parallel,            heap_setup, heap_teardown },
     { "heap/free_routes_to_owner",     test_free_routes_to_owner_list,   heap_setup, heap_teardown },
+    { "heap/reclaim_worker_drains_owner", test_reclaim_worker_drains_owner_list, heap_setup, heap_teardown },
     { "heap/cross_heap_cycle_bounds",  test_cross_heap_cycle_bounds_pools, heap_setup, heap_teardown },
     { "heap/abandon_is_adopted",       test_heap_abandon_is_adopted,     heap_setup, heap_teardown },
     { "heap/alloc_copy_list",          test_alloc_copy_list_retains,     heap_setup, heap_teardown },

test/test_pool.c

@@ -397,6 +397,72 @@ static void pool_count_fn(void* ctx, uint32_t worker_id, int64_t start, int64_t
     }
 }
 
+/* --------------------------------------------------------------------------
+ * Test: blocks a worker allocated inside a dispatch and the main thread freed
+ * after it are back with their owners by the end of the next dispatch — no
+ * registered heap carries a foreign list across a parallel region (issue
+ * #619: a stream of parallel joins grew the process by a pool per worker
+ * while the freed per-task buffers waited on those lists).
+ * -------------------------------------------------------------------------- */
+
+typedef struct {
+    ray_t*   blocks[64];
+    uint32_t owner[64];      /* worker_id that allocated blocks[i]; 0 = main */
+} pool_alloc_ctx_t;
+
+static void pool_alloc_fn(void* ctx, uint32_t worker_id, int64_t start, int64_t end) {
+    pool_alloc_ctx_t* c = (pool_alloc_ctx_t*)ctx;
+    for (int64_t i = start; i < end && i < 64; i++) {
+        c->blocks[i] = ray_alloc(256u << 10);    /* 256 KB from the caller's heap */
+        c->owner[i]  = worker_id;
+    }
+}
+
+static test_result_t test_dispatch_reclaims_worker_blocks(void) {
+    ray_heap_init();
+
+    ray_pool_t pool;
+    ray_err_t err = ray_pool_create(&pool, 3);
+    TEST_ASSERT_EQ_I(err, RAY_OK);
+
+    /* ray_pool_create returns before the workers have started; a worker
+     * publishes its heap once it has run ray_heap_init.  Wait for all three
+     * on that published state, so every slot below is a real heap. */
+    for (uint32_t w = 0; w < pool.n_workers; w++)
+        while (!atomic_load(&pool.worker_heaps[w])) RAY_CPU_RELAX();
+
+    /* Rounds of "workers allocate, main frees" until a round in which at
+     * least one block really came from a worker (main is worker 0 and can
+     * take every task when the others are slow to wake).  Rounds are cheap;
+     * 64 of them without a worker allocation would mean the pool is not
+     * running its workers at all. */
+    pool_alloc_ctx_t ctx = {0};
+    int worker_blocks = 0;
+    for (int round = 0; round < 64 && worker_blocks == 0; round++) {
+        ray_pool_dispatch_n(&pool, pool_alloc_fn, &ctx, 64);
+        for (int i = 0; i < 64; i++) {
+            TEST_ASSERT_NOT_NULL(ctx.blocks[i]);
+            if (ctx.owner[i] != 0) worker_blocks++;
+            ray_free(ctx.blocks[i]);             /* cross-thread for worker blocks */
+            ctx.blocks[i] = NULL;
+        }
+    }
+    TEST_ASSERT(worker_blocks > 0, "some blocks were allocated by workers");
+
+    /* Those frees happened after the last dispatch ended, so the blocks sit
+     * on their owners' foreign lists now; the next dispatch hands them back. */
+    pool_count_ctx_t cctx = {0};
+    ray_pool_dispatch_n(&pool, pool_count_fn, &cctx, 4);
+    for (uint32_t w = 0; w < pool.n_workers; w++) {
+        ray_heap_t* wh = (ray_heap_t*)atomic_load(&pool.worker_heaps[w]);
+        TEST_ASSERT_NOT_NULL(wh);
+        TEST_ASSERT_NULL(atomic_load(&wh->foreign));
+    }
+
+    ray_pool_free(&pool);
+    PASS();
+}
+
 /* --------------------------------------------------------------------------
  * Test: dispatch with total_elems <= 0 returns immediately, no calls fire
  * -------------------------------------------------------------------------- */
@@ -1400,6 +1466,7 @@ const test_entry_t pool_entries[] = {
     { "pool/auto_all_logical_cpus", test_auto_all_logical_cpus, NULL, NULL },
 #endif
     { "pool/parallel_sum", test_parallel_sum, NULL, NULL },
+    { "pool/dispatch_reclaims_worker_blocks", test_dispatch_reclaims_worker_blocks, NULL, NULL },
     { "pool/parallel_add", test_parallel_add, NULL, NULL },
     { "pool/parallel_group_sum", test_parallel_group_sum, NULL, NULL },
     { "pool/parallel_min_max", test_parallel_min_max, NULL, NULL },