Ensure that when a task is evicted, its deps also evicts it from its callers (#644)

This PR adds a fix to `EvictWithCleanup` to ensure that when
a task is evicted, its dependencies also deletes the task from its
callers, otherwise tasks could be holding onto stale references in
its callers list to tasks that have been evicted otherwise.

This PR also adds a debug configuration for `EvictWithCleanup` to
force GC and log keys for keys that should have been evicted when
built in debug mode.

`task.parents` has also been renamed to `task.callers` for clarity.

Author: doriable

experimental/incremental/executor.go

@@ -21,10 +21,13 @@ import (
 	"slices"
 	"sync"
 	"sync/atomic"
+	"time"
+	"weak"
 
 	"golang.org/x/sync/semaphore"
 
 	"github.com/bufbuild/protocompile/experimental/report"
+	"github.com/bufbuild/protocompile/internal"
 )
 
 // Executor is a caching executor for incremental queries.
@@ -42,6 +45,10 @@ type Executor struct {
 
 	sema *semaphore.Weighted
 
+	// The [time.Duration] to wait before running the GC when debug mode is on. See docs for
+	// [WithDebugEvict].
+	evictGCDeadline time.Duration
+
 	counter atomic.Uint64 // Used for generating sequence IDs for Result.Unchanged.
 }
 
@@ -73,6 +80,17 @@ func WithReportOptions(options report.Options) ExecutorOption {
 	return func(e *Executor) { e.reportOptions = options }
 }
 
+// WithDebugEvict takes a [time.Duration] and configures debug mode for evictions in the
+// executor.
+//
+// If set and the compiler is built with the debug tag, when [Executor.EvictWithCleanup]
+// is called, all evicted keys will be tracked. Then after eviction, a goroutine will be
+// kicked off to sleep for the configured duration, force a GC run, and then print out all
+// pointers that should be evicted but have not been GC'd.
+func WithDebugEvict(wait time.Duration) ExecutorOption {
+	return func(e *Executor) { e.evictGCDeadline = wait }
+}
+
 // Keys returns a snapshot of the keys of which queries are present (and
 // memoized) in an Executor.
 //
@@ -208,17 +226,46 @@ func (e *Executor) EvictWithCleanup(keys []any, cleanup func()) {
 
 	e.dirty.Lock()
 	defer e.dirty.Unlock()
+
+	var evicted []weak.Pointer[task]
+	logEvictionDebug := internal.Debug && e.evictGCDeadline > 0
 	for n := len(tasks); n > 0; n = len(tasks) {
 		next := tasks[n-1]
 		tasks = tasks[:n-1]
 
-		next.parents.Range(func(k, _ any) bool {
+		for k := range next.callers.Range {
 			tasks = append(tasks, k.(*task)) //nolint:errcheck
-			return true
-		})
+		}
 
 		// Remove the task from the map. Syncronized by the dirty lock.
-		e.tasks.Delete(next.query.Key())
+		t, _ := e.tasks.LoadAndDelete(next.query.Key())
+		if logEvictionDebug {
+			evicted = append(evicted, weak.Make(t.(*task))) //nolint:errcheck
+		}
+
+		// Remove the task from the callers of its deps.
+		for k := range next.deps.Range {
+			k.(*task).callers.Delete(next) //nolint:errcheck
+		}
+	}
+
+	if evicted != nil {
+		go func() {
+			time.Sleep(e.evictGCDeadline)
+			runtime.GC()
+			evicted = slices.DeleteFunc(evicted, func(e weak.Pointer[task]) bool {
+				return e.Value() == nil
+			})
+			for _, e := range evicted {
+				internal.DebugLog(
+					[]any{"exec %p", e},
+					"EvictWithCleanup",
+					"failed to GC evicted task %p: %v",
+					e.Value(),
+					e.Value().query.Key(),
+				)
+			}
+		}()
 	}
 
 	if cleanup != nil {

experimental/incremental/task.go

@@ -222,12 +222,12 @@ func Resolve[T any](caller *Task, queries ...Query[T]) (results []Result[T], exp
 		deps[i] = dep
 
 		// Update dependency graph.
-		parent := caller.task
-		if parent == nil {
+		callerTask := caller.task
+		if callerTask == nil {
 			continue // Root.
 		}
-		parent.deps.Store(dep, struct{}{})
-		dep.parents.Store(parent, struct{}{})
+		callerTask.deps.Store(dep, struct{}{})
+		dep.callers.Store(callerTask, struct{}{})
 	}
 
 	// Schedule all but the first query to run asynchronously.
@@ -294,7 +294,7 @@ type task struct {
 	// Inverse of deps. Contains all tasks that directly depend on this task.
 	// Written by multiple tasks concurrently.
 	// TODO: See the comment on Executor.tasks.
-	parents sync.Map // [*task]struct{}
+	callers sync.Map // [*task]struct{}
 
 	// If this task has not been started yet, this is nil.
 	// Otherwise, if it is complete, result.done will be closed.