From 7e9ab756106b9f31ff1996cc805c014bf296b4c7 Mon Sep 17 00:00:00 2001 From: Yevgeniy Miretskiy Date: Fri, 17 Feb 2023 13:09:50 -0500 Subject: [PATCH] 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 | 4 +++ pkg/util/pool/BUILD.bazel | 22 +++++++++++++ pkg/util/pool/pool.go | 53 +++++++++++++++++++++++++++++++ pkg/util/pool/pool_test.go | 64 ++++++++++++++++++++++++++++++++++++++ 4 files changed, 143 insertions(+) create mode 100644 pkg/util/pool/BUILD.bazel create mode 100644 pkg/util/pool/pool.go create mode 100644 pkg/util/pool/pool_test.go diff --git a/pkg/BUILD.bazel b/pkg/BUILD.bazel index 59366dd588c3..841c183e95ad 100644 --- a/pkg/BUILD.bazel +++ b/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", diff --git a/pkg/util/pool/BUILD.bazel b/pkg/util/pool/BUILD.bazel new file mode 100644 index 000000000000..786cddcc3f2d --- /dev/null +++ b/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") diff --git a/pkg/util/pool/pool.go b/pkg/util/pool/pool.go new file mode 100644 index 000000000000..09edfad95bb8 --- /dev/null +++ b/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) +} diff --git a/pkg/util/pool/pool_test.go b/pkg/util/pool/pool_test.go new file mode 100644 index 000000000000..58837610ad0e --- /dev/null +++ b/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()) +}