feat(sync): add Pool

sync.go

@@ -47,7 +47,7 @@ type WaitGroup struct {
 }
 
 // GoN starts several concurrent tasks and increases the internal count by len(fs).
-// The internal count will be descreased by one when each of the task is done.
+// The internal count will be decreased by one when each of the task is done.
 // Note: if a call to any function in fs panics, then the whole program crashes.
 //
 // See: https://github.com/golang/go/issues/18022 and https://github.com/golang/go/issues/76126
@@ -75,7 +75,7 @@ func (wg *WaitGroup) Go(fs ...func()) {
 }
 
 // GoN starts a task n times concurrently and increases the internal count by n.
-// The internal count will be descreased by one when each of the task instances is done.
+// The internal count will be decreased by one when each of the task instances is done.
 // Note: if the f call panics, then the whole program crashes.
 //
 // See: https://github.com/golang/go/issues/18022 and https://github.com/golang/go/issues/76126
@@ -108,7 +108,7 @@ func (wg *WaitGroup) Wait() {
 
 // WaitChannel returns a channel which reads will block until the internal counter is zero.
 func (wg *WaitGroup) WaitChannel() <-chan struct{} {
-	var c = make(chan struct{})
+	c := make(chan struct{})
 
 	go func() {
 		wg.wg.Wait()
@@ -117,3 +117,30 @@ func (wg *WaitGroup) WaitChannel() <-chan struct{} {
 
 	return c
 }
+
+type Pool[T any] struct {
+	pool  sync.Pool
+	New   func() T
+	Reset func(T) T
+}
+
+func (p *Pool[T]) Put(x T) {
+	p.pool.Put(x)
+}
+
+func (p *Pool[T]) Get() (x T) {
+	v := p.pool.Get()
+	if v == nil {
+		if p.New != nil {
+			return p.New()
+		}
+		return x
+	}
+
+	x = v.(T)
+	if p.Reset != nil {
+		x = p.Reset(x)
+	}
+
+	return x
+}

sync_test.go

@@ -1,12 +1,15 @@
 package nstd
 
 import (
+	"math"
+	"runtime"
+	"runtime/debug"
 	"testing"
 )
 
 func TestMutexAndWaitGroup(t *testing.T) {
 	const N = 1000
-	var n = 0
+	n := 0
 	defer func() {
 		if expected := N * 12; n != expected {
 			t.Fatalf("n != %d", expected)
@@ -15,7 +18,7 @@ func TestMutexAndWaitGroup(t *testing.T) {
 
 	var wg WaitGroup
 
-	var f = func() {
+	f := func() {
 		var m Mutex
 		for range [1000]struct{}{} {
 			wg.Go(func() {
@@ -38,3 +41,98 @@ func TestMutexAndWaitGroup(t *testing.T) {
 	defer wg.Wait()
 	f()
 }
+
+func TestPool(t *testing.T) {
+	// disable GC so we can control when it happens.
+	defer debug.SetGCPercent(debug.SetGCPercent(-1))
+	p := Pool[*rune]{}
+	if p.Get() != nil {
+		t.Fatal("expected empty")
+	}
+
+	runes := [...]rune{'a', 'b', 'c'}
+
+	pinner := new(runtime.Pinner)
+	pinner.Pin(&p)
+	p.Put(&runes[0])
+	p.Put(&runes[1])
+	if g := p.Get(); g == nil || *g != runes[0] {
+		t.Fatalf("got %#v; want %c", g, runes[0])
+	}
+	if g := p.Get(); g == nil || *g != runes[1] {
+		t.Fatalf("got %#v; want %c", g, runes[1])
+	}
+	if g := p.Get(); g != nil {
+		t.Fatalf("got %#v; want nil", g)
+	}
+	pinner.Unpin()
+
+	// Put in a large number of objects so they spill into
+	// stealable space.
+	for range 100 {
+		p.Put(&runes[2])
+	}
+	// After one GC, the victim cache should keep them alive.
+	runtime.GC()
+	if g := p.Get(); g == nil || *g != runes[2] {
+		t.Fatalf("got %#v; want %c after GC", g, runes[2])
+	}
+	// A second GC should drop the victim cache.
+	runtime.GC()
+	if g := p.Get(); g != nil {
+		t.Fatalf("got %#v; want nil after second GC", g)
+	}
+}
+
+func TestPoolNew(t *testing.T) {
+	// disable GC so we can control when it happens.
+	defer debug.SetGCPercent(debug.SetGCPercent(-1))
+
+	i := 0
+	p := Pool[int]{
+		New: func() int {
+			i++
+			return i
+		},
+	}
+	if v := p.Get(); v != 1 {
+		t.Fatalf("got %v; want 1", v)
+	}
+	if v := p.Get(); v != 2 {
+		t.Fatalf("got %v; want 2", v)
+	}
+
+	pinner := new(runtime.Pinner)
+	pinner.Pin(&p)
+	p.Put(42)
+	if v := p.Get(); v != 42 {
+		t.Fatalf("got %v; want 42", v)
+	}
+	pinner.Unpin()
+
+	if v := p.Get(); v != 3 {
+		t.Fatalf("got %v; want 3", v)
+	}
+}
+
+func TestPoolReset(t *testing.T) {
+	// disable GC so we can control when it happens.
+	defer debug.SetGCPercent(debug.SetGCPercent(-1))
+
+	to := math.MinInt
+	p := Pool[*int]{
+		Reset: func(v *int) *int {
+			*v = to
+			return v
+		},
+	}
+
+	pinner := new(runtime.Pinner)
+	pinner.Pin(&p)
+	x := 42
+	p.Put(&x)
+	if v := p.Get(); v == nil || *v != to {
+		t.Fatalf("got %v; want %d", v, to)
+	}
+	pinner.Unpin()
+}