diff --git a/sync.go b/sync.go index e374ee1..f69ad32 100644 --- a/sync.go +++ b/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 +} diff --git a/sync_test.go b/sync_test.go index b6c0d78..81ff816 100644 --- a/sync_test.go +++ b/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() +}