Tidy up and unify sync pool handling (#8068)

anderseknert · web-flow · commit dd2ccc70249b · 2025-11-21T10:52:22.000+01:00
As @srenatus pointed out in a previous review, these were a little messy in some places. This moves the generic SyncPool type to the util package for reuse across the codebase. Also some various related improvements. Signed-off-by: Anders Eknert <anders.eknert@apple.com>
diff --git a/v1/ast/syncpools.go b/v1/ast/syncpools.go
@@ -4,12 +4,15 @@ import (
 	"bytes"
 	"strings"
 	"sync"
+
+	"github.com/open-policy-agent/opa/v1/util"
 )
 
 var (
-	TermPtrPool     = NewSyncPool[Term]()
-	BytesReaderPool = NewSyncPool[bytes.Reader]()
-	IndexResultPool = NewSyncPool[IndexResult]()
+	TermPtrPool     = util.NewSyncPool[Term]()
+	BytesReaderPool = util.NewSyncPool[bytes.Reader]()
+	IndexResultPool = util.NewSyncPool[IndexResult]()
+	bbPool          = util.NewSyncPool[bytes.Buffer]()
 	// Needs custom pool because of custom Put logic.
 	sbPool = &stringBuilderPool{
 		pool: sync.Pool{
@@ -29,37 +32,10 @@ var (
 )
 
 type (
-	syncPool[T any] struct {
-		pool sync.Pool
-	}
-	stringBuilderPool struct {
-		pool sync.Pool
-	}
-	vvPool struct {
-		pool sync.Pool
-	}
+	stringBuilderPool struct{ pool sync.Pool }
+	vvPool            struct{ pool sync.Pool }
 )
 
-func NewSyncPool[T any]() *syncPool[T] {
-	return &syncPool[T]{
-		pool: sync.Pool{
-			New: func() any {
-				return new(T)
-			},
-		},
-	}
-}
-
-func (p *syncPool[T]) Get() *T {
-	return p.pool.Get().(*T)
-}
-
-func (p *syncPool[T]) Put(x *T) {
-	if x != nil {
-		p.pool.Put(x)
-	}
-}
-
 func (p *stringBuilderPool) Get() *strings.Builder {
 	return p.pool.Get().(*strings.Builder)
 }
diff --git a/v1/ast/term.go b/v1/ast/term.go
@@ -1153,12 +1153,6 @@ func IsVarCompatibleString(s string) bool {
 	return varRegexp.MatchString(s)
 }
 
-var bbPool = &sync.Pool{
-	New: func() any {
-		return new(bytes.Buffer)
-	},
-}
-
 func (ref Ref) String() string {
 	// Note(anderseknert):
 	// Options tried in the order of cheapness, where after some effort,
@@ -1180,7 +1174,7 @@ func (ref Ref) String() string {
 
 	_var := ref[0].Value.String()
 
-	bb := bbPool.Get().(*bytes.Buffer)
+	bb := bbPool.Get()
 	bb.Reset()
 
 	defer bbPool.Put(bb)
diff --git a/v1/topdown/eval.go b/v1/topdown/eval.go
@@ -120,25 +120,10 @@ type eval struct {
 }
 
 type (
-	evp struct {
-		pool sync.Pool
-	}
-	evfp struct {
-		pool sync.Pool
-	}
-	evbp struct {
-		pool sync.Pool
-	}
+	evfp struct{ pool sync.Pool }
+	evbp struct{ pool sync.Pool }
 )
 
-func (ep *evp) Put(e *eval) {
-	ep.pool.Put(e)
-}
-
-func (ep *evp) Get() *eval {
-	return ep.pool.Get().(*eval)
-}
-
 func (ep *evfp) Put(e *evalFunc) {
 	if e != nil {
 		e.e, e.terms, e.ir = nil, nil, nil
@@ -162,13 +147,9 @@ func (ep *evbp) Get() *evalBuiltin {
 }
 
 var (
-	evalPool = &evp{
-		pool: sync.Pool{
-			New: func() any {
-				return &eval{}
-			},
-		},
-	}
+	evalPool     = util.NewSyncPool[eval]()
+	deecPool     = util.NewSyncPool[deferredEarlyExitContainer]()
+	resolverPool = util.NewSyncPool[evalResolver]()
 	evalFuncPool = &evfp{
 		pool: sync.Pool{
 			New: func() any {
@@ -1688,7 +1669,7 @@ func (e *eval) getRules(ref ast.Ref, args []*ast.Term) (*ast.IndexResult, error)
 		return nil, nil
 	}
 
-	resolver := resolverPool.Get().(*evalResolver)
+	resolver := resolverPool.Get()
 	defer func() {
 		resolver.e = nil
 		resolver.args = nil
@@ -1743,14 +1724,6 @@ type evalResolver struct {
 	args []*ast.Term
 }
 
-var (
-	resolverPool = sync.Pool{
-		New: func() any {
-			return &evalResolver{}
-		},
-	}
-)
-
 func (e *evalResolver) Resolve(ref ast.Ref) (ast.Value, error) {
 	e.e.instr.startTimer(evalOpResolve)
 
@@ -2439,12 +2412,6 @@ func (dc *deferredEarlyExitContainer) copyError() *deferredEarlyExitError {
 	return &cpy
 }
 
-var deecPool = sync.Pool{
-	New: func() any {
-		return &deferredEarlyExitContainer{}
-	},
-}
-
 type evalTree struct {
 	e         *eval
 	bindings  *bindings
@@ -2530,7 +2497,7 @@ func (e evalTree) enumerate(iter unifyIterator) error {
 		return err
 	}
 
-	dc := deecPool.Get().(*deferredEarlyExitContainer)
+	dc := deecPool.Get()
 	dc.deferred = nil
 	defer deecPool.Put(dc)
 
diff --git a/v1/util/performance.go b/v1/util/performance.go
@@ -7,6 +7,34 @@ import (
 	"unsafe"
 )
 
+// SyncPool is a generic sync.Pool for type T, providing some convenience
+// over sync.Pool directly: [SyncPool.Put] ensures that nil values are not
+// put into the pool, and [SyncPool.Get] returns a pointer to T without having
+// to do a type assertion at the call site.
+type SyncPool[T any] struct {
+	pool sync.Pool
+}
+
+func NewSyncPool[T any]() *SyncPool[T] {
+	return &SyncPool[T]{
+		pool: sync.Pool{
+			New: func() any {
+				return new(T)
+			},
+		},
+	}
+}
+
+func (p *SyncPool[T]) Get() *T {
+	return p.pool.Get().(*T)
+}
+
+func (p *SyncPool[T]) Put(x *T) {
+	if x != nil {
+		p.pool.Put(x)
+	}
+}
+
 // NewPtrSlice returns a slice of pointers to T with length n,
 // with only 2 allocations performed no matter the size of n.
 // See:
@@ -133,5 +161,7 @@ func (sp *SlicePool[T]) Get(length int) *[]T {
 
 // Put returns a pointer to a slice of type T to the pool.
 func (sp *SlicePool[T]) Put(s *[]T) {
-	sp.pool.Put(s)
+	if s != nil {
+		sp.pool.Put(s)
+	}
 }
diff --git a/v1/util/read_gzip_body.go b/v1/util/read_gzip_body.go
@@ -8,17 +8,11 @@ import (
 	"io"
 	"net/http"
 	"strings"
-	"sync"
 
 	"github.com/open-policy-agent/opa/v1/util/decoding"
 )
 
-var gzipReaderPool = sync.Pool{
-	New: func() any {
-		reader := new(gzip.Reader)
-		return reader
-	},
-}
+var gzipReaderPool = NewSyncPool[gzip.Reader]()
 
 // Note(philipc): Originally taken from server/server.go
 // The DecodingLimitHandler handles validating that the gzip payload is within the
@@ -49,7 +43,7 @@ func ReadMaybeCompressedBody(r *http.Request) ([]byte, error) {
 			return nil, errors.New("gzip payload too large")
 		}
 
-		gzReader := gzipReaderPool.Get().(*gzip.Reader)
+		gzReader := gzipReaderPool.Get()
 		defer func() {
 			gzReader.Close()
 			gzipReaderPool.Put(gzReader)
