test: StateDB.StopPrefetcher() blocks on WorkerPool.Wait()

Author: ARR4N

core/state/trie_prefetcher.go

@@ -79,6 +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()
 	for _, fetcher := range p.fetchers {
 		fetcher.abort() // safe to do multiple times
 
@@ -303,9 +304,11 @@ func (sf *subfetcher) abort() {
 // loop waits for new tasks to be scheduled and keeps loading them until it runs
 // out of tasks or its underlying trie is retrieved for committing.
 func (sf *subfetcher) loop() {
-	defer sf.pool.wait()
 	// No matter how the loop stops, signal anyone waiting that it's terminated
-	defer close(sf.term)
+	defer func() {
+		sf.pool.wait()
+		close(sf.term)
+	}()
 
 	// Start by opening the trie and stop processing if it fails
 	if sf.owner == (common.Hash{}) {

core/state/trie_prefetcher.libevm.go

@@ -68,6 +68,21 @@ 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() {
+	var wg sync.WaitGroup
+	for _, f := range p.fetchers {
+		wg.Add(1)
+		go func(f *subfetcher) {
+			f.abort()
+			wg.Done()
+		}(f)
+	}
+	wg.Wait()
+}
+
 func (p *subfetcherPool) wait() {
 	if p == nil || p.workers == nil {
 		return

core/state/trie_prefetcher.libevm_test.go

@@ -0,0 +1,74 @@
+// Copyright 2024 the libevm authors.
+//
+// The libevm additions to go-ethereum are free software: you can redistribute
+// them and/or modify them under the terms of the GNU Lesser General Public License
+// as published by the Free Software Foundation, either version 3 of the License,
+// or (at your option) any later version.
+//
+// The libevm additions are distributed in the hope that they will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser
+// General Public License for more details.
+//
+// You should have received a copy of the GNU Lesser General Public License
+// along with the go-ethereum library. If not, see
+// <http://www.gnu.org/licenses/>.
+
+package state
+
+import (
+	"testing"
+	"time"
+
+	"github.com/ava-labs/libevm/common"
+)
+
+type synchronisingWorkerPool struct {
+	executed, unblock chan struct{}
+}
+
+var _ WorkerPool = (*synchronisingWorkerPool)(nil)
+
+func (p *synchronisingWorkerPool) Execute(func()) {
+	select {
+	case <-p.executed:
+	default:
+		close(p.executed)
+	}
+}
+
+func (p *synchronisingWorkerPool) Wait() {
+	<-p.unblock
+}
+
+func TestStopPrefetcherWaitsOnWorkers(t *testing.T) {
+	pool := &synchronisingWorkerPool{
+		executed: make(chan struct{}),
+		unblock:  make(chan struct{}),
+	}
+	opt := WithWorkerPools(func() WorkerPool { return pool })
+
+	db := filledStateDB()
+	db.prefetcher = newTriePrefetcher(db.db, db.originalRoot, "", opt)
+	db.prefetcher.prefetch(common.Hash{}, common.Hash{}, common.Address{}, [][]byte{{}})
+
+	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.
+		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)
+	}
+}