From 8c0982b8857bcf7659cd7d0b76671782377368f2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=BA=90=E6=96=87=E9=9B=A8?= <41315874+fumiama@users.noreply.github.com> Date: Wed, 26 Feb 2025 16:35:01 +0900 Subject: [PATCH] feat(pbuf): add USRDAT --- README.md | 25 ++++++++++++++------ pbuf/buffer.go | 12 +++++++--- pbuf/buffer_test.go | 8 +++---- pbuf/bytes.go | 56 +++++++++++++++++++++++---------------------- pbuf/pbuf.go | 24 ++++++++++++------- pbuf/pooler.go | 24 +++++++++++-------- pool_test.go | 18 ++++++++++++--- 7 files changed, 105 insertions(+), 62 deletions(-) diff --git a/README.md b/README.md index e7da505..b833333 100644 --- a/README.md +++ b/README.md @@ -3,16 +3,27 @@ Lightweight & Safe (buffer-writer | general object) pool. ## Quick Start ```go -import ( - "crypto/rand" +package main - "github.com/fumiama/orbyte/pbuf" +import ( + "crypto/rand" + "io" + + "github.com/fumiama/orbyte/pbuf" ) func main() { - b := pbuf.NewBytes(1024) // Allocate Bytes from pool. - rand.Read(b.Bytes()) // Do sth. - b.KeepAlive() // Mark as reachable. - // After that, b will be auto-reused on GC. + buf := pbuf.NewBuffer(nil) // Allocate Buffer from pool. + buf.P(func(buf *pbuf.Buffer) { + io.CopyN(buf, rand.Reader, 4096) // Do sth. + }) + // After that, buf will be auto-reused on GC. + + b := pbuf.NewBytes(1024) // Allocate Bytes from pool. + b.V(func(b []byte) { + rand.Read(b) // Do sth. + }) + // After that, b will be auto-reused on GC. + } ``` diff --git a/pbuf/buffer.go b/pbuf/buffer.go index f391108..d6006b5 100644 --- a/pbuf/buffer.go +++ b/pbuf/buffer.go @@ -7,16 +7,22 @@ import ( ) // NewBuffer wraps bytes.NewBuffer into Item. -func (bufferPool BufferPool) NewBuffer(buf []byte) *orbyte.Item[bytes.Buffer] { +func (bufferPool BufferPool[USRDAT]) NewBuffer( + buf []byte, +) *orbyte.Item[UserBuffer[USRDAT]] { return bufferPool.p.New(buf) } // InvolveBuffer involve external *bytes.Buffer into Item. -func (bufferPool BufferPool) InvolveBuffer(buf *bytes.Buffer) *orbyte.Item[bytes.Buffer] { +func (bufferPool BufferPool[USRDAT]) InvolveBuffer( + buf *bytes.Buffer, +) *orbyte.Item[UserBuffer[USRDAT]] { return bufferPool.p.Involve(buf.Len(), buf) } // ParseBuffer convert external *bytes.Buffer into Item. -func (bufferPool BufferPool) ParseBuffer(buf *bytes.Buffer) *orbyte.Item[bytes.Buffer] { +func (bufferPool BufferPool[USRDAT]) ParseBuffer( + buf *bytes.Buffer, +) *orbyte.Item[UserBuffer[USRDAT]] { return bufferPool.p.Parse(buf.Len(), buf) } diff --git a/pbuf/buffer_test.go b/pbuf/buffer_test.go index 9d8f139..9cbdb32 100644 --- a/pbuf/buffer_test.go +++ b/pbuf/buffer_test.go @@ -6,8 +6,6 @@ import ( "io" "runtime" "testing" - - "github.com/fumiama/orbyte" ) func TestBuffer(t *testing.T) { @@ -19,8 +17,8 @@ func TestBuffer(t *testing.T) { testBuffer(InvolveBuffer(bytes.NewBuffer(make([]byte, 0, 8192))), t) } -func testBuffer(buf *orbyte.Item[bytes.Buffer], t *testing.T) { - buf.P(func(buf *bytes.Buffer) { +func testBuffer(buf *OBuffer, t *testing.T) { + buf.P(func(buf *Buffer) { if buf.Len() != 4096 { io.CopyN(buf, rand.Reader, 4096) if buf.Len() != 4096 { @@ -30,7 +28,7 @@ func testBuffer(buf *orbyte.Item[bytes.Buffer], t *testing.T) { }) bufcp := buf.Copy() dat := buf.Trans() - bufcp.P(func(bufcp *bytes.Buffer) { + bufcp.P(func(bufcp *Buffer) { if bufcp.Len() != 4096 { t.Fatal("got", bufcp.Len()) } diff --git a/pbuf/bytes.go b/pbuf/bytes.go index cca2245..58c99bc 100644 --- a/pbuf/bytes.go +++ b/pbuf/bytes.go @@ -6,41 +6,43 @@ import ( "github.com/fumiama/orbyte" ) -// Bytes wrap pooled buffer into []byte +// UserBytes wrap pooled buffer into []byte // while sharing the same pool. -type Bytes struct { - buf *orbyte.Item[bytes.Buffer] +type UserBytes[USRDAT any] struct { + buf *orbyte.Item[UserBuffer[USRDAT]] a, b int } -// BufferItemToBytes convert between *orbyte.Item[bytes.Buffer] +// BufferItemToBytes convert between *Buffer // and Bytes. // // Please notice that Bytes cannnot convert back to -// *orbyte.Item[bytes.Buffer] again. -func BufferItemToBytes(buf *orbyte.Item[bytes.Buffer]) (b Bytes) { +// *Buffer again. +func BufferItemToBytes[USRDAT any]( + buf *orbyte.Item[UserBuffer[USRDAT]], +) (b UserBytes[USRDAT]) { b.buf = buf - buf.P(func(buf *bytes.Buffer) { + buf.P(func(buf *UserBuffer[USRDAT]) { b.b = buf.Len() }) return } // NewBytes alloc sz bytes. -func (bufferPool BufferPool) NewBytes(sz int) (b Bytes) { +func (bufferPool BufferPool[USRDAT]) NewBytes(sz int) (b UserBytes[USRDAT]) { buf := bufferPool.p.New(sz) b.buf = buf - buf.P(func(buf *bytes.Buffer) { + buf.P(func(buf *UserBuffer[USRDAT]) { b.b = buf.Len() }) return } // InvolveBytes involve outside buf into pool. -func (bufferPool BufferPool) InvolveBytes(p ...byte) (b Bytes) { +func (bufferPool BufferPool[USRDAT]) InvolveBytes(p ...byte) (b UserBytes[USRDAT]) { buf := bufferPool.p.Involve(len(p), bytes.NewBuffer(p)) b.buf = buf - buf.P(func(buf *bytes.Buffer) { + buf.P(func(buf *UserBuffer[USRDAT]) { b.b = buf.Len() }) return @@ -48,10 +50,10 @@ func (bufferPool BufferPool) InvolveBytes(p ...byte) (b Bytes) { // ParseBytes convert outside bytes to Bytes safely // without adding it into pool. -func (bufferPool BufferPool) ParseBytes(p ...byte) (b Bytes) { +func (bufferPool BufferPool[USRDAT]) ParseBytes(p ...byte) (b UserBytes[USRDAT]) { buf := bufferPool.p.Parse(len(p), bytes.NewBuffer(p)) b.buf = buf - buf.P(func(buf *bytes.Buffer) { + buf.P(func(buf *UserBuffer[USRDAT]) { b.b = buf.Len() }) return @@ -59,54 +61,54 @@ func (bufferPool BufferPool) ParseBytes(p ...byte) (b Bytes) { // HasInit whether this Bytes is made by pool or // just declared. -func (b Bytes) HasInit() bool { +func (b UserBytes[USRDAT]) HasInit() bool { return b.buf != nil } // Trans please refer to Item.Trans(). -func (b Bytes) Trans() []byte { +func (b UserBytes[USRDAT]) Trans() []byte { buf := b.buf.Trans() return buf.Bytes()[b.a:b.b] } // Len of slice. -func (b Bytes) Len() int { +func (b UserBytes[USRDAT]) Len() int { return b.b - b.a } // Cap of slice. -func (b Bytes) Cap() (c int) { - b.buf.P(func(b *bytes.Buffer) { +func (b UserBytes[USRDAT]) Cap() (c int) { + b.buf.P(func(b *UserBuffer[USRDAT]) { c = b.Cap() }) return c } // V use the inner value safely -func (b Bytes) V(f func([]byte)) { - b.buf.P(func(buf *bytes.Buffer) { +func (b UserBytes[USRDAT]) V(f func([]byte)) { + b.buf.P(func(buf *UserBuffer[USRDAT]) { f(buf.Bytes()[b.a:b.b]) }) } // Copy please refer to Item.Copy(). -func (b Bytes) Copy() (cb Bytes) { +func (b UserBytes[USRDAT]) Copy() (cb UserBytes[USRDAT]) { cb.buf = b.buf.Copy() cb.a, cb.b = b.a, b.b return } // SliceFrom dat[from:] with Ref. -func (b Bytes) SliceFrom(from int) Bytes { - return Bytes{buf: b.buf, a: b.a + from, b: b.b} +func (b UserBytes[USRDAT]) SliceFrom(from int) UserBytes[USRDAT] { + return UserBytes[USRDAT]{buf: b.buf, a: b.a + from, b: b.b} } // SliceTo dat[:to] with Ref. -func (b Bytes) SliceTo(to int) Bytes { - return Bytes{buf: b.buf, a: b.a, b: b.a + to} +func (b UserBytes[USRDAT]) SliceTo(to int) UserBytes[USRDAT] { + return UserBytes[USRDAT]{buf: b.buf, a: b.a, b: b.a + to} } // Slice dat[from:to] with Ref. -func (b Bytes) Slice(from, to int) Bytes { - return Bytes{buf: b.buf, a: b.a + from, b: b.a + to} +func (b UserBytes[USRDAT]) Slice(from, to int) UserBytes[USRDAT] { + return UserBytes[USRDAT]{buf: b.buf, a: b.a + from, b: b.a + to} } diff --git a/pbuf/pbuf.go b/pbuf/pbuf.go index 8a6384a..71d296c 100644 --- a/pbuf/pbuf.go +++ b/pbuf/pbuf.go @@ -7,28 +7,36 @@ import ( "github.com/fumiama/orbyte" ) -var bufferPool = NewBufferPool() +var bufferPool = NewBufferPool[struct{}]() -type BufferPool struct { - p *orbyte.Pool[bytes.Buffer] +type ( + OBuffer = orbyte.Item[Buffer] + Buffer = UserBuffer[struct{}] + Bytes = UserBytes[struct{}] +) + +type BufferPool[USRDAT any] struct { + p *orbyte.Pool[UserBuffer[USRDAT]] } -func NewBufferPool() BufferPool { - return BufferPool{p: orbyte.NewPool[bytes.Buffer](bufpooler{})} +func NewBufferPool[USRDAT any]() BufferPool[USRDAT] { + return BufferPool[USRDAT]{ + p: orbyte.NewPool[UserBuffer[USRDAT]](bufpooler[USRDAT]{}), + } } // NewBuffer wraps bytes.NewBuffer into Item. -func NewBuffer(buf []byte) *orbyte.Item[bytes.Buffer] { +func NewBuffer(buf []byte) *OBuffer { return bufferPool.NewBuffer(buf) } // InvolveBuffer involve external *bytes.Buffer into Item. -func InvolveBuffer(buf *bytes.Buffer) *orbyte.Item[bytes.Buffer] { +func InvolveBuffer(buf *bytes.Buffer) *OBuffer { return bufferPool.InvolveBuffer(buf) } // ParseBuffer convert external *bytes.Buffer into Item. -func ParseBuffer(buf *bytes.Buffer) *orbyte.Item[bytes.Buffer] { +func ParseBuffer(buf *bytes.Buffer) *OBuffer { return bufferPool.ParseBuffer(buf) } diff --git a/pbuf/pooler.go b/pbuf/pooler.go index 4df4d25..0ffe7d4 100644 --- a/pbuf/pooler.go +++ b/pbuf/pooler.go @@ -7,9 +7,15 @@ import ( "unsafe" ) -type bufpooler struct{} +// UserBuffer with customizable user data structure inside. +type UserBuffer[USRDAT any] struct { + bytes.Buffer + DAT USRDAT +} -func (bufpooler) New(config any, pooled bytes.Buffer) bytes.Buffer { +type bufpooler[USRDAT any] struct{} + +func (bufpooler[USRDAT]) New(config any, pooled UserBuffer[USRDAT]) UserBuffer[USRDAT] { switch c := config.(type) { case int: pooled.Grow(c) @@ -24,7 +30,7 @@ func (bufpooler) New(config any, pooled bytes.Buffer) bytes.Buffer { if len(c) != buf.Len() { panic("unexpected bad bytes.NewBuffer") } - return *buf + return UserBuffer[USRDAT]{Buffer: *buf} } return pooled case string: @@ -35,12 +41,12 @@ func (bufpooler) New(config any, pooled bytes.Buffer) bytes.Buffer { } } -func (bufpooler) Parse(obj any, pooled bytes.Buffer) bytes.Buffer { +func (bufpooler[USRDAT]) Parse(obj any, pooled UserBuffer[USRDAT]) UserBuffer[USRDAT] { switch o := obj.(type) { case *bytes.Buffer: - return *o + return UserBuffer[USRDAT]{Buffer: *o} case bytes.Buffer: - return o + return UserBuffer[USRDAT]{Buffer: o} case []byte: pooled.Write(o) return pooled @@ -58,17 +64,17 @@ func (bufpooler) Parse(obj any, pooled bytes.Buffer) bytes.Buffer { } } -func (bufpooler) Reset(item *bytes.Buffer) { +func (bufpooler[USRDAT]) Reset(item *UserBuffer[USRDAT]) { // See https://golang.org/issue/23199 const maxSize = 1 << 16 if item.Cap() > maxSize { // drop large buffer - *item = bytes.Buffer{} + *item = UserBuffer[USRDAT]{} return } item.Reset() } -func (bufpooler) Copy(dst, src *bytes.Buffer) { +func (bufpooler[USRDAT]) Copy(dst, src *UserBuffer[USRDAT]) { dst.Reset() srccp := *src _, err := io.Copy(dst, &srccp) diff --git a/pool_test.go b/pool_test.go index d91c597..de46594 100644 --- a/pool_test.go +++ b/pool_test.go @@ -1,7 +1,9 @@ package orbyte import ( + "bytes" "crypto/rand" + "encoding/hex" "runtime" "sync" "testing" @@ -29,13 +31,21 @@ func TestPool(t *testing.T) { if out != 0 { t.Fatal("unexpected behavior") } + item := p.New(4096) + item.V(func(b []byte) { + rand.Read(b) + }) + exp := item.Copy().Trans() wg := sync.WaitGroup{} for i := 0; i < 4096; i++ { item := p.New(i) + item.V(func(b []byte) { + copy(b, exp) + }) wg.Add(1) go useranddes(item.Copy(), &wg) wg.Add(1) - go userv(item.Trans(), &wg) + go userv(item.Trans(), &wg, exp[:i]) } wg.Wait() runtime.GC() @@ -54,9 +64,11 @@ func useranddes(item *Item[[]byte], wg *sync.WaitGroup) { item.ManualDestroy() } -func userv(b []byte, wg *sync.WaitGroup) { +func userv(b []byte, wg *sync.WaitGroup, exp []byte) { defer wg.Done() - rand.Read(b) + if !bytes.Equal(b, exp) { + panic("expect " + hex.EncodeToString(exp) + " got " + hex.EncodeToString(b)) + } } type simplepooler struct{}