refactor: subfetchers track outstanding tasks

Author: ARR4N

core/state/trie_prefetcher.go

@@ -79,7 +79,7 @@ func newTriePrefetcher(db Database, root common.Hash, namespace string, opts ...
 // close iterates over all the subfetchers, aborts any that were left spinning
 // and reports the stats to the metrics subsystem.
 func (p *triePrefetcher) close() {
-	p.abortFetchersConcurrently()
+	p.abortFetchersAndReleaseWorkerPools()
 	for _, fetcher := range p.fetchers {
 		fetcher.abort() // safe to do multiple times
 

core/state/trie_prefetcher.libevm.go

@@ -31,10 +31,12 @@ type prefetcherConfig struct {
 	newWorkers func() WorkerPool
 }
 
-// A WorkerPool executes functions asynchronously.
+// A WorkerPool executes functions asynchronously. Done() is called to signal
+// that the pool is no longer needed and that Execute() is guaranteed to not be
+// called again.
 type WorkerPool interface {
 	Execute(func())
-	Wait()
+	Done()
 }
 
 // WithWorkerPools configures trie prefetching to execute asynchronously. The
@@ -49,6 +51,7 @@ func WithWorkerPools(ctor func() WorkerPool) PrefetcherOption {
 type subfetcherPool struct {
 	workers WorkerPool
 	tries   sync.Pool
+	wg      sync.WaitGroup
 }
 
 // applyTo configures the [subfetcher] to use a [WorkerPool] if one was provided
@@ -68,10 +71,9 @@ func (c *prefetcherConfig) applyTo(sf *subfetcher) {
 	}
 }
 
-// abortFetchersConcurrently calls [subfetcher.abort] on every fetcher, blocking
-// until all return. Calling abort() sequentially may result in later fetchers
-// accepting new work in the interim.
-func (p *triePrefetcher) abortFetchersConcurrently() {
+func (p *triePrefetcher) abortFetchersAndReleaseWorkerPools() {
+	// Calling abort() sequentially may result in later fetchers accepting new
+	// work in the interim.
 	var wg sync.WaitGroup
 	for _, f := range p.fetchers {
 		wg.Add(1)
@@ -80,26 +82,35 @@ func (p *triePrefetcher) abortFetchersConcurrently() {
 			wg.Done()
 		}(f)
 	}
+
+	// A WorkerPool is allowed to be shared between fetchers so we MUST wait for
+	// them to finish all tasks otherwise they could call Execute() after
+	// Done(), which we guarantee in the public API to be impossible.
 	wg.Wait()
+	for _, f := range p.fetchers {
+		if w := f.pool.workers; w != nil {
+			w.Done()
+		}
+	}
 }
 
 func (p *subfetcherPool) wait() {
-	if p == nil || p.workers == nil {
-		return
-	}
-	p.workers.Wait()
+	p.wg.Wait()
 }
 
 // execute runs the provided function with a copy of the subfetcher's Trie.
 // Copies are stored in a [sync.Pool] to reduce creation overhead. If p was
 // configured with a [WorkerPool] then it is used for function execution,
 // otherwise `fn` is just called directly.
 func (p *subfetcherPool) execute(fn func(Trie)) {
+	p.wg.Add(1)
 	do := func() {
 		t := p.tries.Get().(Trie)
 		fn(t)
 		p.tries.Put(t)
+		p.wg.Done()
 	}
+
 	if w := p.workers; w != nil {
 		w.Execute(do)
 	} else {

core/state/trie_prefetcher.libevm_test.go

@@ -21,28 +21,39 @@ import (
 	"time"
 
 	"github.com/ava-labs/libevm/common"
+	"github.com/stretchr/testify/assert"
 )
 
 type synchronisingWorkerPool struct {
-	executed, unblock chan struct{}
+	t                             *testing.T
+	executed, unblock             chan struct{}
+	done                          bool
+	preconditionsToStopPrefetcher int
 }
 
 var _ WorkerPool = (*synchronisingWorkerPool)(nil)
 
-func (p *synchronisingWorkerPool) Execute(func()) {
+func (p *synchronisingWorkerPool) Execute(fn func()) {
+	fn()
 	select {
 	case <-p.executed:
 	default:
 		close(p.executed)
 	}
-}
 
-func (p *synchronisingWorkerPool) Wait() {
 	<-p.unblock
+	assert.False(p.t, p.done, "Done() called before Execute() returns")
+	p.preconditionsToStopPrefetcher++
+}
+
+func (p *synchronisingWorkerPool) Done() {
+	p.done = true
+	p.preconditionsToStopPrefetcher++
 }
 
 func TestStopPrefetcherWaitsOnWorkers(t *testing.T) {
 	pool := &synchronisingWorkerPool{
+		t:        t,
 		executed: make(chan struct{}),
 		unblock:  make(chan struct{}),
 	}
@@ -55,20 +66,14 @@ func TestStopPrefetcherWaitsOnWorkers(t *testing.T) {
 	go func() {
 		<-pool.executed
 		// Sleep otherwise there is a small chance that we close pool.unblock
-		// between db.StopPrefetcher() returning and the select receiving on the
-		// channel.
+		// between db.StopPrefetcher() returning and the assertion.
 		time.Sleep(time.Second)
 		close(pool.unblock)
 	}()
 
 	<-pool.executed
 	db.StopPrefetcher()
-	select {
-	case <-pool.unblock:
-		// The channel was closed, therefore pool.Wait() unblocked. This is a
-		// necessary pre-condition for db.StopPrefetcher() unblocking, and the
-		// purpose of this test.
-	default:
-		t.Errorf("%T.StopPrefetcher() returned before %T.Wait() unblocked", db, pool)
-	}
+	// If this happens then either Execute() hadn't returned or Done() wasn't
+	// called.
+	assert.Equalf(t, 2, pool.preconditionsToStopPrefetcher, "%T.StopPrefetcher() returned early", db)
 }