diff --git a/tokenizer/lattice/mem/doc.go b/tokenizer/lattice/mem/doc.go index 2213700..d9e28c8 100644 --- a/tokenizer/lattice/mem/doc.go +++ b/tokenizer/lattice/mem/doc.go @@ -1,2 +1,2 @@ -// Package mem implements the memory utility. +// Package mem implements the memory utility such as memory pool. package mem diff --git a/tokenizer/lattice/mem/pool.go b/tokenizer/lattice/mem/pool.go index 1366ed0..6c2f8e4 100644 --- a/tokenizer/lattice/mem/pool.go +++ b/tokenizer/lattice/mem/pool.go @@ -5,16 +5,25 @@ import ( ) // Pool represents memory pool of T. +// +// It is suitable for managing temporary objects that can be individually saved +// and retrieved (mem.Pool.Put and mem.Pool.Get). +// Unlike variables or pointer variables, mem.Pool is safe for use by multiple +// goroutines and no new memory will be allocated. +// +// It is a wrapper of sync.Pool to support generics and easy to use. For the +// actual usage, see the example in the package documentation. type Pool[T any] struct { internal *sync.Pool } -// NewPool returns a memory pool of T. -func NewPool[T any](f func() *T) Pool[T] { +// NewPool returns a memory pool of T. newF is a constructor of T, and it is called +// when the memory pool is empty. +func NewPool[T any](newF func() *T) Pool[T] { return Pool[T]{ internal: &sync.Pool{ New: func() any { - return f() + return newF() }, }, } diff --git a/tokenizer/lattice/mem/pool_test.go b/tokenizer/lattice/mem/pool_test.go new file mode 100644 index 0000000..7e50587 --- /dev/null +++ b/tokenizer/lattice/mem/pool_test.go @@ -0,0 +1,119 @@ +package mem_test + +import ( + "fmt" + "sync" + "testing" + + "github.com/ikawaha/kagome/v2/tokenizer/lattice/mem" +) + +func ExamplePool() { + type Foo struct { + Bar string + } + + // newFoo is a constructor of Foo type. + newFoo := func() *Foo { + return &Foo{ + Bar: "let the foo begin", + } + } + + // Create a new memory pool of Foo type. + // If the memory pool is empty, it creates a new instance of Foo using newFoo. + bufPool := mem.NewPool[Foo](newFoo) + + // Retrieve a Foo instance from the memory pool and print the current value + // of the Bar field. + a := bufPool.Get() + fmt.Println(a.Bar) + + // Set the Bar field then put it back to the memory pool. + a.Bar = "buz" + bufPool.Put(a) + + // Same as above but set a different value to the Bar field. + // + // This will overwrite the previous value. But note that this will not allocate + // new memory and is safe for use by multiple goroutines simultaneously. + // See the benchmark in the same test file. + b := bufPool.Get() + b.Bar = "qux" + bufPool.Put(b) + + // Retrieve a Foo instance from the memory pool and print the current value + // of the Bar field. + c := bufPool.Get() + fmt.Println(c.Bar) + // Output: + // let the foo begin + // qux +} + +// To benchmark run: +// +// go test -benchmem -bench=Benchmark -count 5 ./tokenizer/lattice/mem +func BenchmarkPool(b *testing.B) { + type Foo struct { + Bar int + } + + bufPool := mem.NewPool[Foo](func() *Foo { + return new(Foo) + }) + + b.ResetTimer() + + // Spawn 3 goroutines to get and put the Foo instance from the memory pool + // b.N times each. + // + // Note that mem.Pool is safe for use by multiple goroutines simultaneously + // and no new memory will be allocated. + var wg sync.WaitGroup + + for i := 0; i < 3; i++ { + wg.Add(1) + + go func(max int) { + defer wg.Done() + + for i := 0; i < max; i++ { + a := bufPool.Get() // get + a.Bar += i // increment + bufPool.Put(a) // put + } + }(b.N) + } + + wg.Wait() +} + +// To fuzz test run: +// +// go test -fuzz=Fuzz -fuzztime 1m ./tokenizer/lattice/mem/... +func FuzzPool(f *testing.F) { + type Foo struct { + Bar string + } + + bufPool := mem.NewPool[Foo](func() *Foo { + return new(Foo) + }) + + // Corpus variants/patterns to fuzz test. + f.Add("") + f.Add(" ") + f.Add("0123456789") + f.Add("!@#$%^&*()_+") + f.Add("short") + f.Add("短") + f.Add("this is a test with a long string") + f.Add("これは、いささか長いテスト用の文字列です。") + + f.Fuzz(func(_ *testing.T, s string) { + a := bufPool.Get() // get + a.Bar += s // append + bufPool.Put(a) // put + }) +}