diff --git a/src/core/pool.c b/src/core/pool.c index c3cac800..904600fd 100644 --- a/src/core/pool.c +++ b/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); } /* -------------------------------------------------------------------------- diff --git a/src/core/pool.h b/src/core/pool.h index db0ea49f..f9914df0 100644 --- a/src/core/pool.h +++ b/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 diff --git a/src/mem/heap.c b/src/mem/heap.c index 906ef225..45244056 100644 --- a/src/mem/heap.c +++ b/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) * diff --git a/src/mem/heap.h b/src/mem/heap.h index 9b8b6129..84a25828 100644 --- a/src/mem/heap.h +++ b/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); diff --git a/test/test_heap.c b/test/test_heap.c index 0e14473f..975eb7f6 100644 --- a/test/test_heap.c +++ b/test/test_heap.c @@ -41,6 +41,7 @@ #include "test.h" #include #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 }, diff --git a/test/test_pool.c b/test/test_pool.c index 8eb499a1..6a3069f1 100644 --- a/test/test_pool.c +++ b/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 },