fix(pool): release the last job's arrays; no writes after a recovered caller panic

SECURITY.md

@@ -57,6 +57,24 @@ standard library's known vulnerabilities; `govulncheck ./...` reports none.
 
 ## Audit log
 
+**2026-10-07 (v0.6.1).** An audit of the helper pool added in v0.6.0 found
+two defects. Both are fixed, each with a test that fails without the fix.
+
+- **Retention.** The pool kept the last job, whose closure holds the
+  operation's arrays, so the arrays of the last parallel operation could not
+  be collected until another one replaced it: up to the size of the largest
+  array, for as long as the program ran. `TestPoolDoesNotRetainLastJob` checks
+  the data is collectable once `parallelFor` returns. A first version of that
+  test passed on the defective code, because its closure captured a variable
+  the test then set to nil.
+- **Writes after a recovered panic.** Since v0.6.0 the caller runs blocks
+  itself. If one panicked there, for example a `Map` function, the panic
+  propagated while the helpers went on writing. A caller that recovered could
+  find its arrays changing under it. The caller now stops handing out blocks
+  and waits for the running ones before the panic propagates
+  (`TestPoolCallerPanicStopsTheHelpers`). A panic on a helper goroutine still
+  ends the program, as it did before v0.6.0; `Map`'s documentation says so.
+
 **2026-10-06 (v0.5.1).**
 
 - The Go floor is 1.27.1 since v0.5.0. `govulncheck` reports no vulnerability

internal/kernels/parallel.go

@@ -62,24 +62,45 @@ func parallelFor(n, w int, body func(lo, hi int)) {
 	start := func(i int) int { return i*q + min(i, r) }
 	own := func(g int) int { return g * blocks / w } // owner g: [own(g), own(g+1))
 	taken := make([]paddedCounter, w)
-	var done atomic.Int64
+	var done, running atomic.Int64
 	run := func(g int) {
 		for k := 0; k < w; k++ {
 			o := (g + k) % w
 			lo, hi := own(o), own(o+1)
 			for {
+				running.Add(1)
 				i := lo + int(taken[o].Add(1)) - 1
 				if i >= hi {
+					running.Add(-1)
 					break
 				}
 				body(start(i), start(i+1))
 				done.Add(1)
+				running.Add(-1)
 			}
 		}
 	}
 	publish(run, w)
+	// If body panics here, in the caller, stop handing out blocks and wait for
+	// the ones the helpers are running before the panic propagates: a caller
+	// that recovers must not see its arrays written to afterwards. Closing the
+	// runs matters even though the wait alone usually outlasts the work:
+	// running can touch zero for an instant while every helper is between two
+	// blocks, and without the closing more blocks could start after that.
+	defer func() {
+		if p := recover(); p != nil {
+			for o := range taken {
+				taken[o].Store(int64(blocks) + 1)
+			}
+			running.Add(-1) // the caller's own block, the one that panicked
+			spinUntil(func() bool { return running.Load() == 0 })
+			retire()
+			panic(p)
+		}
+	}()
 	run(0)
 	spinUntil(func() bool { return done.Load() == int64(blocks) })
+	retire()
 }
 
 // paddedCounter is an atomic counter alone on its cache line, so owners

internal/kernels/pool.go

@@ -84,6 +84,15 @@ func helperLoop(i int, h *poolHelper, seen uint64) {
 	}
 }
 
+// noJob replaces a finished job in pool.job. Left there, the job's closure would
+// keep the operation's arrays alive until the next parallel operation. Its w of
+// 0 admits no helper.
+var noJob = &poolJob{}
+
+// retire drops the reference to the finished job. A helper that loaded it just
+// before finds every block taken and lets go of it on return.
+func retire() { pool.job.Store(noJob) }
+
 // publish hands run to the first w-1 helpers and wakes the parked ones.
 func publish(run func(g int), w int) {
 	pool.job.Store(&poolJob{run: run, w: w})

internal/kernels/pool_retention_test.go

@@ -0,0 +1,64 @@
+package kernels
+
+import (
+	"runtime"
+	"sync/atomic"
+	"testing"
+	"time"
+	"weak"
+)
+
+// TestPoolDoesNotRetainLastJob: once parallelFor returns, nothing in the pool
+// may keep the job's data alive. The last job used to stay in pool.job, its
+// closure holding the operands, so the arrays of the last parallel operation
+// could not be collected until another one replaced it.
+func TestPoolDoesNotRetainLastJob(t *testing.T) {
+	big := make([]float64, 1<<20)
+	ref := weak.Make(&big[0])
+	fill(big) // the closure captures fill's parameter, which nothing clears
+	big = nil
+	for i := 0; i < 5 && ref.Value() != nil; i++ {
+		runtime.GC()
+	}
+	if ref.Value() != nil {
+		t.Fatal("the last parallelFor job keeps its data alive after returning")
+	}
+}
+
+func fill(a []float64) {
+	parallelFor(len(a), 4, func(lo, hi int) {
+		for i := lo; i < hi; i++ {
+			a[i] = 1
+		}
+	})
+}
+
+// TestPoolCallerPanicStopsTheHelpers: when body panics in the caller, the
+// panic propagates only after the helpers have stopped, and no block starts
+// afterwards; the pool then serves the next call.
+func TestPoolCallerPanicStopsTheHelpers(t *testing.T) {
+	var after atomic.Int64
+	var stopped atomic.Bool
+	func() {
+		defer func() {
+			if recover() == nil {
+				t.Fatal("the panic did not propagate")
+			}
+			stopped.Store(true)
+		}()
+		parallelFor(64, 4, func(lo, hi int) {
+			if lo == 0 {
+				panic("boom") // block 0 is the first the caller claims
+			}
+			time.Sleep(time.Millisecond)
+			if stopped.Load() {
+				after.Add(1)
+			}
+		})
+	}()
+	time.Sleep(20 * time.Millisecond)
+	if n := after.Load(); n != 0 {
+		t.Fatalf("%d blocks ran after the panic propagated", n)
+	}
+	checkCover(t, 1000, 4)
+}

ndarray.go

@@ -750,7 +750,10 @@ func (a *Array) DivScalar(v float64) *Array { r, _ := a.Div(scalarArray(v)); ret
 // Map returns a new contiguous array with f applied to every element. For a
 // contiguous receiver the elements are read in place (no materialise copy); the
 // elementwise pass is parallelised across cores above the kernel threshold. f
-// must be safe to call concurrently (the package's math ufuncs are).
+// must be safe to call concurrently (the package's math ufuncs are), and should
+// not panic: on a large array a panic may happen on a worker goroutine, which
+// ends the program as any unrecovered goroutine panic does. A panic on the
+// calling goroutine propagates once the workers have stopped writing.
 func (a *Array) Map(f func(float64) float64) *Array {
 	src := a.contiguousData()
 	dst := a.alloc(len(src), false)