util/pool: Add generic implementation of sync.Pool

sync.Pool is used to ammortize memory allocations.
When using sync.Pool it is important, for efficiency reasons,
to store objects of the same type and, preferably.

This makes it hard to use sync.Pool over generic types.
this PR adds `pool.Pool[T]` abstraction for managing
objects of type T.

Epic: None
Release note: None

pkg/BUILD.bazel

@@ -631,6 +631,7 @@ ALL_TESTS = [
     "//pkg/util/netutil/addr:addr_test",
     "//pkg/util/netutil:netutil_test",
     "//pkg/util/optional:optional_test",
+    "//pkg/util/pool:pool_test",
     "//pkg/util/pprofutil:pprofutil_test",
     "//pkg/util/pretty:pretty_test",
     "//pkg/util/protoutil:protoutil_test",
@@ -2212,6 +2213,8 @@ GO_TARGETS = [
     "//pkg/util/netutil:netutil_test",
     "//pkg/util/optional:optional",
     "//pkg/util/optional:optional_test",
+    "//pkg/util/pool:pool",
+    "//pkg/util/pool:pool_test",
     "//pkg/util/pprofutil:pprofutil",
     "//pkg/util/pprofutil:pprofutil_test",
     "//pkg/util/pretty:pretty",
@@ -3304,6 +3307,7 @@ GET_X_DATA_TARGETS = [
     "//pkg/util/netutil:get_x_data",
     "//pkg/util/netutil/addr:get_x_data",
     "//pkg/util/optional:get_x_data",
+    "//pkg/util/pool:get_x_data",
     "//pkg/util/pprofutil:get_x_data",
     "//pkg/util/pretty:get_x_data",
     "//pkg/util/protoutil:get_x_data",

pkg/util/pool/BUILD.bazel

@@ -0,0 +1,22 @@
+load("//build/bazelutil/unused_checker:unused.bzl", "get_x_data")
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
+
+go_library(
+    name = "pool",
+    srcs = ["pool.go"],
+    importpath = "github.com/cockroachdb/cockroach/pkg/util/pool",
+    visibility = ["//visibility:public"],
+)
+
+go_test(
+    name = "pool_test",
+    srcs = ["pool_test.go"],
+    args = ["-test.timeout=295s"],
+    deps = [
+        ":pool",
+        "//pkg/util/leaktest",
+        "@com_github_stretchr_testify//require",
+    ],
+)
+
+get_x_data(name = "get_x_data")

pkg/util/pool/pool.go

@@ -0,0 +1,53 @@
+// Copyright 2023 The Cockroach Authors.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+//
+
+package pool
+
+import (
+	"reflect"
+	"sync"
+)
+
+// Pool is a generic implementation of sync.Pool.
+// Normally, sync.Pool stores objects of the same type.
+// This implementation adopts sync.Pool to type T, and
+// reuses a single underlying sync.Pool per type.
+type Pool[T any] struct {
+	*sync.Pool
+	initFn func(*T)
+}
+
+var typedPools sync.Map
+
+// MakePool creates Pool[T] which returns *T, and initializes those
+// pointers as per init function.
+func MakePool[T any](initFn func(*T)) Pool[T] {
+	var zero T
+	typ := reflect.TypeOf(zero)
+
+	p, ok := typedPools.Load(typ)
+	if !ok {
+		p, _ = typedPools.LoadOrStore(typ, &sync.Pool{New: func() any { return new(T) }})
+	}
+	return Pool[T]{Pool: p.(*sync.Pool), initFn: initFn}
+}
+
+// Get returns initialized *T.
+func (p *Pool[T]) Get() *T {
+	v := p.Pool.Get().(*T)
+	p.initFn(v)
+	return v
+}
+
+// Put returns value to this pool.
+func (p *Pool[T]) Put(v *T) {
+	p.Pool.Put(v)
+}

pkg/util/pool/pool_test.go

@@ -0,0 +1,64 @@
+// Copyright 2023 The Cockroach Authors.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+//
+
+package pool_test
+
+import (
+	"testing"
+
+	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
+	"github.com/cockroachdb/cockroach/pkg/util/pool"
+	"github.com/stretchr/testify/require"
+)
+
+func TestTypedPool(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+
+	intPool := pool.MakePool[int](func(i *int) {
+		*i = 0
+	})
+
+	intPool2 := pool.MakePool[int](func(i *int) {
+		*i = 42
+	})
+
+	type rec struct {
+		a, b int
+	}
+
+	type rec2 struct {
+		a, b int
+	}
+	recPool := pool.MakePool[rec](func(r *rec) {
+		r.a = 1
+		r.b = -1
+	})
+	recPool2 := pool.MakePool[rec2](func(r *rec2) {
+		r.a = 1
+		r.b = -1
+	})
+
+	// Underlying sync.Pool are reused per type.
+	require.Equal(t, intPool.Pool, intPool2.Pool)
+	require.NotEqual(t, recPool.Pool, recPool2.Pool)
+	require.NotEqual(t, intPool.Pool, recPool.Pool)
+
+	// returned objects are initialized as per init function.
+	a := intPool.Get()
+	b := intPool2.Get()
+	require.NotEqual(t, a, b) // a, b are different objects.
+	require.Equal(t, 0, *a)
+	require.Equal(t, 42, *b)
+	intPool.Put(a)
+	intPool2.Put(b)
+
+	require.Equal(t, &rec{a: 1, b: -1}, recPool.Get())
+}