feat: libevm/sync.Pool[T] for type safety

Author: ARR4N

core/state/trie_prefetcher.libevm.go

@@ -17,10 +17,9 @@
 package state
 
 import (
-	"sync"
-
 	"github.com/ava-labs/libevm/common"
 	"github.com/ava-labs/libevm/libevm/options"
+	"github.com/ava-labs/libevm/libevm/sync"
 	"github.com/ava-labs/libevm/log"
 )
 
@@ -50,18 +49,18 @@ func WithWorkerPools(ctor func() WorkerPool) PrefetcherOption {
 
 type subfetcherPool struct {
 	workers WorkerPool
-	tries   sync.Pool
+	tries   sync.Pool[Trie]
 	wg      sync.WaitGroup
 }
 
 // applyTo configures the [subfetcher] to use a [WorkerPool] if one was provided
 // with a [PrefetcherOption].
 func (c *prefetcherConfig) applyTo(sf *subfetcher) {
 	sf.pool = &subfetcherPool{
-		tries: sync.Pool{
+		tries: sync.Pool[Trie]{
 			// Although the workers may be shared between all subfetchers, each
 			// MUST have its own Trie pool.
-			New: func() any {
+			New: func() Trie {
 				return sf.db.CopyTrie(sf.trie)
 			},
 		},
@@ -94,7 +93,7 @@ func (p *subfetcherPool) wait() {
 func (p *subfetcherPool) execute(fn func(Trie)) {
 	p.wg.Add(1)
 	do := func() {
-		t := p.tries.Get().(Trie)
+		t := p.tries.Get()
 		fn(t)
 		p.tries.Put(t)
 		p.wg.Done()

core/state/trie_prefetcher.libevm_test.go

@@ -20,8 +20,9 @@ import (
 	"testing"
 	"time"
 
-	"github.com/ava-labs/libevm/common"
 	"github.com/stretchr/testify/assert"
+
+	"github.com/ava-labs/libevm/common"
 )
 
 type synchronisingWorkerPool struct {
@@ -73,7 +74,7 @@ func TestStopPrefetcherWaitsOnWorkers(t *testing.T) {
 
 	<-pool.executed
 	db.StopPrefetcher()
-	// If this happens then either Execute() hadn't returned or Done() wasn't
+	// If this errors then either Execute() hadn't returned or Done() wasn't
 	// called.
 	assert.Equalf(t, 2, pool.preconditionsToStopPrefetcher, "%T.StopPrefetcher() returned early", db)
 }

libevm/sync/sync.go

@@ -0,0 +1,52 @@
+// 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 sync extends the standard library's sync package.
+package sync
+
+import "sync"
+
+// Aliases of stdlib sync's types to avoid having to import it alongside this
+// package.
+type (
+	Cond      = sync.Cond
+	Locker    = sync.Locker
+	Map       = sync.Map
+	Mutex     = sync.Mutex
+	Once      = sync.Once
+	RWMutex   = sync.RWMutex
+	WaitGroup = sync.WaitGroup
+)
+
+// A Pool is a type-safe wrapper around [sync.Pool].
+type Pool[T any] struct {
+	New  func() T
+	pool sync.Pool
+	once Once
+}
+
+// Get is equivalent to [sync.Pool.Get].
+func (p *Pool[T]) Get() T {
+	p.once.Do(func() { // Do() guarantees at least once, not just only once
+		p.pool.New = func() any { return p.New() }
+	})
+	return p.pool.Get().(T) //nolint:forcetypeassert
+}
+
+// Put is equivalent to [sync.Pool.Put].
+func (p *Pool[T]) Put(t T) {
+	p.pool.Put(t)
+}