From 61ab79747ce83bde7cedc642ab261849247e5193 Mon Sep 17 00:00:00 2001
From: Stephen Dolan <sdolan@janestreet.com>
Date: Tue, 28 Jan 2025 11:47:53 +0000
Subject: [PATCH] Ensure that pool owners are correctly set on pool adoption

(Plus a new compaction test for the failure triggered by getting this wrong)
---
 runtime/shared_heap.c                         |  4 ++++
 .../compaction/test_compact_manydomains.ml    | 21 +++++++++++++++++++
 2 files changed, 25 insertions(+)
 create mode 100644 testsuite/tests/compaction/test_compact_manydomains.ml

diff --git a/runtime/shared_heap.c b/runtime/shared_heap.c
index c818cd2f24b..936dd413e29 100644
--- a/runtime/shared_heap.c
+++ b/runtime/shared_heap.c
@@ -366,6 +366,7 @@ static pool* pool_global_adopt(struct caml_heap_state* local, sizeclass sz)
     if( r ) {
       atomic_store_relaxed(&pool_freelist.global_avail_pools[sz], r->next);
       r->next = 0;
+      r->owner = local->owner;
       local->avail_pools[sz] = r;
       adopt_pool_stats_with_lock(local, r, sz);
 
@@ -390,6 +391,7 @@ static pool* pool_global_adopt(struct caml_heap_state* local, sizeclass sz)
     if( r ) {
       atomic_store_relaxed(&pool_freelist.global_full_pools[sz], r->next);
       r->next = local->full_pools[sz];
+      r->owner = local->owner;
       local->full_pools[sz] = r;
       adopt_pool_stats_with_lock(local, r, sz);
 
@@ -405,6 +407,7 @@ static pool* pool_global_adopt(struct caml_heap_state* local, sizeclass sz)
       pool_sweep(local, &local->full_pools[sz], sz, 0);
     r = local->avail_pools[sz];
   }
+  CAMLassert(r == NULL || r->owner == local->owner);
   return r;
 }
 
@@ -574,6 +577,7 @@ static intnat pool_sweep(struct caml_heap_state* local, pool** plist,
     header_t* end = POOL_END(a);
     mlsize_t wh = wsize_sizeclass[sz];
     int all_used = 1;
+    CAMLassert(a->owner == local->owner);
 
     /* conceptually, this is incremented by [wh] for every iteration
        below, however we can hoist these increments knowing that [p ==
diff --git a/testsuite/tests/compaction/test_compact_manydomains.ml b/testsuite/tests/compaction/test_compact_manydomains.ml
new file mode 100644
index 00000000000..29b94b6a48e
--- /dev/null
+++ b/testsuite/tests/compaction/test_compact_manydomains.ml
@@ -0,0 +1,21 @@
+(* TEST
+ flags += "-alert -unsafe_parallelism";
+ runtime5;
+ { bytecode; }
+ { native; }
+*)
+
+let num_domains = 20
+
+let go () =
+  let n = 50_000 in
+  let c = Array.make n None in
+  for i = 0 to n-1 do
+    c.(i) <- Some (i, i)
+  done;
+  Gc.compact ()
+
+let () =
+  Array.init num_domains (fun _ -> Domain.spawn go)
+  |> Array.iter Domain.join
+