DArray/stencil: Automatically free HaloArray cache

Author: jpsamaroo

src/array/stencil.jl

@@ -729,15 +729,16 @@ function select_neighborhood_info(chunks, idx, neigh_dist, boundary)
     return region_metadata, neighbor_chunks
 end
 
-# Per-thread cache: IdDict{DArray, Dict{(chunk_idx, halo_width), HaloArray}}.
-# Using IdDict for the outer level ensures two DArrays with identical element types and
-# chunk shapes never share a buffer. Using chunk_idx as part of the inner key ensures that
-# within one DArray, every chunk has its own dedicated buffer — so if a single worker thread
-# processes multiple same-shaped chunks in the same iteration (sequentially), each gets a
-# distinct HaloArray and there is no aliasing with a concurrently running inner-stencil task.
+# Per-thread cache: WeakKeyDict{DArray, Dict{(chunk_idx, halo_width), HaloArray}}.
+# WeakKeyDict is used for the outer level so that the cache does not hold a strong reference
+# to the source DArray — allowing its GC finalizer to fire when user code drops its last
+# reference (see below). Using chunk_idx as part of the inner key ensures that within one
+# DArray, every chunk has its own dedicated buffer — so if a single worker thread processes
+# multiple same-shaped chunks in the same iteration sequentially, each gets a distinct
+# HaloArray and there is no aliasing with a concurrently running inner-stencil task.
 # Filling a cached buffer in-place is safe because spawn_datadeps blocks until all inner
 # tasks complete before the next iteration's build_halo_consolidated calls run.
-const HALO_ARRAY_CACHE = TaskLocalValue{IdDict{Any,Dict{Any,Any}}}(()->IdDict{Any,Dict{Any,Any}}())
+const HALO_ARRAY_CACHE = TaskLocalValue{WeakKeyDict{Any,Dict{Any,Any}}}(()->WeakKeyDict{Any,Dict{Any,Any}}())
 
 # Consolidated halo builder: loads all neighbor regions directly into a HaloArray.
 # `read_darray` and `chunk_idx` are used solely for cache lookup — they are not DTask
@@ -752,7 +753,22 @@ function build_halo_consolidated(read_darray, chunk_idx, neigh_dist, boundary, c
     halo_width = ntuple(i -> get_neigh_dist(neigh_dist, i), N)
 
     outer_cache = HALO_ARRAY_CACHE[]
-    inner_cache = get!(outer_cache, read_darray) do; Dict{Any,Any}(); end
+
+    # Create the inner cache on first encounter of this DArray on this thread, and register
+    # a finalizer that captures it. When the DArray becomes unreachable and is collected,
+    # the finalizer fires and unsafe_free!s every cached HaloArray for this (DArray, thread)
+    # pair. Because WeakKeyDict holds only a weak reference to read_darray, the DArray can
+    # actually be collected (a plain IdDict would keep it alive forever).
+    if !haskey(outer_cache, read_darray)
+        inner_cache = Dict{Any,Any}()
+        outer_cache[read_darray] = inner_cache
+        finalizer(read_darray) do _
+            for halo in values(inner_cache)
+                unsafe_free!(halo)
+            end
+        end
+    end
+    inner_cache = outer_cache[read_darray]
     cache_key = (chunk_idx, halo_width)
 
     if haskey(inner_cache, cache_key)