From e92f4c6af85dd72654b45821700f845001dcd3b1 Mon Sep 17 00:00:00 2001
From: Joe Tsai <joetsai@digital-static.net>
Date: Wed, 10 Jul 2024 09:39:52 -0700
Subject: [PATCH] syncs: add generic Pool (#12759)

Pool is a type-safe wrapper over sync.Pool.

Updates tailscale/corp#11038
Updates #cleanup

Signed-off-by: Joe Tsai <joetsai@digital-static.net>
---
 syncs/pool.go      | 31 +++++++++++++++++++++++++++++++
 syncs/pool_test.go | 30 ++++++++++++++++++++++++++++++
 2 files changed, 61 insertions(+)
 create mode 100644 syncs/pool.go
 create mode 100644 syncs/pool_test.go

diff --git a/syncs/pool.go b/syncs/pool.go
new file mode 100644
index 000000000..46ffd2e52
--- /dev/null
+++ b/syncs/pool.go
@@ -0,0 +1,31 @@
+// Copyright (c) Tailscale Inc & AUTHORS
+// SPDX-License-Identifier: BSD-3-Clause
+
+package syncs
+
+import "sync"
+
+// Pool is the generic version of [sync.Pool].
+type Pool[T any] struct {
+	pool sync.Pool
+
+	// New optionally specifies a function to generate
+	// a value when Get would otherwise return the zero value of T.
+	// It may not be changed concurrently with calls to Get.
+	New func() T
+}
+
+// Get selects an arbitrary item from the Pool, removes it from the Pool,
+// and returns it to the caller. See [sync.Pool.Get].
+func (p *Pool[T]) Get() T {
+	x, ok := p.pool.Get().(T)
+	if !ok && p.New != nil {
+		x = p.New()
+	}
+	return x
+}
+
+// Put adds x to the pool.
+func (p *Pool[T]) Put(x T) {
+	p.pool.Put(x)
+}
diff --git a/syncs/pool_test.go b/syncs/pool_test.go
new file mode 100644
index 000000000..798b18cba
--- /dev/null
+++ b/syncs/pool_test.go
@@ -0,0 +1,30 @@
+// Copyright (c) Tailscale Inc & AUTHORS
+// SPDX-License-Identifier: BSD-3-Clause
+
+package syncs
+
+import "testing"
+
+func TestPool(t *testing.T) {
+	var pool Pool[string]
+	s := pool.Get() // should not panic
+	if s != "" {
+		t.Fatalf("got %q, want %q", s, "")
+	}
+	pool.New = func() string { return "new" }
+	s = pool.Get()
+	if s != "new" {
+		t.Fatalf("got %q, want %q", s, "new")
+	}
+	var found bool
+	for range 1000 {
+		pool.Put("something")
+		found = pool.Get() == "something"
+		if found {
+			break
+		}
+	}
+	if !found {
+		t.Fatalf("unable to get any value put in the pool")
+	}
+}