Skip to content

Commit

Permalink
Merge pull request #324 from KEINOS/add-tests-for-memPool
Browse files Browse the repository at this point in the history
chore: add tests and examples for mem.Pool
  • Loading branch information
ikawaha authored Apr 20, 2024
2 parents 9cf6609 + c08f94e commit 08386d5
Show file tree
Hide file tree
Showing 3 changed files with 130 additions and 2 deletions.
2 changes: 1 addition & 1 deletion tokenizer/lattice/mem/doc.go
Original file line number Diff line number Diff line change
@@ -1,2 +1,2 @@
// Package mem implements the memory utility.
// Package mem implements the memory utility such as memory pool.
package mem
11 changes: 10 additions & 1 deletion tokenizer/lattice/mem/pool.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,11 +5,20 @@ 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.
// NewPool returns a memory pool of T. f is a constructor of T and it is called
// when the memory pool is empty.
func NewPool[T any](f func() *T) Pool[T] {
return Pool[T]{
internal: &sync.Pool{
Expand Down
119 changes: 119 additions & 0 deletions tokenizer/lattice/mem/pool_test.go
Original file line number Diff line number Diff line change
@@ -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
})
}

0 comments on commit 08386d5

Please sign in to comment.