Merge pull request #3 from anantadwi13/dev

reduce memory allocation

Author: anantadwi13

io/buffer.go

@@ -147,7 +147,7 @@ type bufReader struct {
 	isClosed         int32
 	isEofReached     bool
 	reader           io.ReadCloser
-	buffer           [][]byte
+	buffer           []*Buffer
 
 	currentPos int64
 }
@@ -282,7 +282,7 @@ func (b *bufReader) readTo(p []byte) (n int, err error) {
 		buf := b.buffer[b.currentPos/int64(b.pool.BufferSize())]
 		currentPos := int(b.currentPos % int64(b.pool.BufferSize()))
 
-		read := copy(p[n:], buf[currentPos:])
+		read := copy(p[n:], buf.buffer[currentPos:])
 		n += read
 		b.currentPos += int64(read)
 	}
@@ -308,12 +308,12 @@ func (b *bufReader) read(n int64) (bytesRead int64, err error) {
 			return
 		}
 
-		var buf []byte
+		var buf *Buffer
 
 		if len(b.buffer) != 0 {
 			buf = b.buffer[len(b.buffer)-1]
 		}
-		if buf == nil || len(buf) == cap(buf) {
+		if buf == nil || len(buf.buffer) == cap(buf.buffer) {
 			buf, err = b.pool.Get(b.ctx)
 			if err != nil {
 				if errors.Is(err, context.Canceled) {
@@ -322,14 +322,14 @@ func (b *bufReader) read(n int64) (bytesRead int64, err error) {
 				return
 			}
 
-			buf = buf[:0]
+			buf.buffer = buf.buffer[:0]
 			b.buffer = append(b.buffer, buf)
 		}
 
 		var tmpN int
-		tmpN, err = b.reader.Read(buf[len(buf):cap(buf)])
+		tmpN, err = b.reader.Read(buf.buffer[len(buf.buffer):cap(buf.buffer)])
 		if tmpN > 0 {
-			b.buffer[len(b.buffer)-1] = b.buffer[len(b.buffer)-1][:len(buf)+tmpN]
+			buf.buffer = buf.buffer[:len(buf.buffer)+tmpN]
 			bytesRead += int64(tmpN)
 		}
 	}
@@ -342,7 +342,7 @@ func (b *bufReader) getReaderPos() int64 {
 		return 0
 	}
 
-	return int64(l-1)*int64(b.pool.BufferSize()) + int64(len(b.buffer[l-1]))
+	return int64(l-1)*int64(b.pool.BufferSize()) + int64(len(b.buffer[l-1].buffer))
 }
 
 func (b *bufReader) cleanUpBuffer(all bool) {
@@ -356,7 +356,7 @@ func (b *bufReader) cleanUpBuffer(all bool) {
 		if b.buffer[i] == nil {
 			continue
 		}
-		b.pool.Put(b.buffer[i])
+		b.buffer[i].cleanUp()
 		b.buffer[i] = nil
 	}
 	b.buffer = nil

io/buffer_test.go

@@ -147,7 +147,7 @@ func TestReaderVariant(t *testing.T) {
 }
 
 func TestFlowNormalRead(t *testing.T) {
-	tp := &testPool{bufSize: 5}
+	tp := &testPool{p: newPool(5)}
 	bf := NewBufferReadSeekCloserFactory(OptionWithPool(tp))
 	assert.EqualValues(t, 5, bf.BufferSize())
 
@@ -211,7 +211,7 @@ func TestFlowNormalRead(t *testing.T) {
 }
 
 func TestFlowReadSeek(t *testing.T) {
-	tp := &testPool{bufSize: 5}
+	tp := &testPool{p: newPool(5)}
 	bf := NewBufferReadSeekCloserFactory(OptionWithPool(tp))
 	assert.EqualValues(t, 5, bf.BufferSize())
 
@@ -327,7 +327,7 @@ func TestFlowReadSeek(t *testing.T) {
 }
 
 func TestFlowReadSeekOutOfRange(t *testing.T) {
-	tp := &testPool{bufSize: 5}
+	tp := &testPool{p: newPool(5)}
 	bf := NewBufferReadSeekCloserFactory(OptionWithPool(tp))
 	assert.EqualValues(t, 5, bf.BufferSize())
 
@@ -362,7 +362,7 @@ func TestFlowReadSeekOutOfRange(t *testing.T) {
 }
 
 func TestFlowReadSeekDisableSeeker(t *testing.T) {
-	tp := &testPool{bufSize: 5}
+	tp := &testPool{p: newPool(5)}
 	bf := NewBufferReadSeekCloserFactory(OptionWithPool(tp))
 	assert.EqualValues(t, 5, bf.BufferSize())
 
@@ -418,7 +418,7 @@ func TestFlowReadSeekDisableSeeker(t *testing.T) {
 }
 
 func TestFlowWithBufferNotFullyFilled(t *testing.T) {
-	tp := &testPool{bufSize: 5}
+	tp := &testPool{p: newPool(5)}
 	bf := NewBufferReadSeekCloserFactory(OptionWithPool(tp))
 	assert.EqualValues(t, 5, bf.BufferSize())
 

io/types.go

@@ -29,8 +29,26 @@ type BufferReadSeekCloser interface {
 	DisableSeeker()
 }
 
+type Buffer struct {
+	pool   Pool
+	buffer []byte
+}
+
+func NewBuffer(pool Pool, buf []byte) *Buffer {
+	buf = buf[:cap(buf)]
+	return &Buffer{
+		pool:   pool,
+		buffer: buf,
+	}
+}
+
+func (b *Buffer) cleanUp() {
+	b.buffer = b.buffer[:cap(b.buffer)]
+	b.pool.Put(b)
+}
+
 type Pool interface {
 	BufferSize() int
-	Put(buf []byte)
-	Get(ctx context.Context) ([]byte, error)
+	Put(buf *Buffer)
+	Get(ctx context.Context) (*Buffer, error)
 }

io/util_test.go

@@ -44,31 +44,37 @@ func (p *noPool) BufferSize() int {
 	return p.bufSize
 }
 
-func (p *noPool) Put(buf []byte) {
+func (p *noPool) Put(buf *Buffer) {
 }
 
-func (p *noPool) Get(ctx context.Context) ([]byte, error) {
-	return make([]byte, p.bufSize), nil
+func (p *noPool) Get(ctx context.Context) (*Buffer, error) {
+	return NewBuffer(p, make([]byte, p.bufSize)), nil
 }
 
 type testPool struct {
-	bufSize int
-	diff    int32
+	diff int32
+	p    Pool
 }
 
 func (t *testPool) Diff() int32 {
 	return atomic.LoadInt32(&t.diff)
 }
 
 func (t *testPool) BufferSize() int {
-	return t.bufSize
+	return t.p.BufferSize()
 }
 
-func (t *testPool) Put(buf []byte) {
+func (t *testPool) Put(buf *Buffer) {
 	atomic.AddInt32(&t.diff, -1)
+	t.p.Put(buf)
 }
 
-func (t *testPool) Get(ctx context.Context) ([]byte, error) {
+func (t *testPool) Get(ctx context.Context) (*Buffer, error) {
 	atomic.AddInt32(&t.diff, 1)
-	return make([]byte, 0, t.bufSize), nil
+	buf, err := t.p.Get(ctx)
+	if err != nil {
+		return nil, err
+	}
+	buf.pool = t
+	return buf, nil
 }

io/utils.go

@@ -20,20 +20,23 @@ func newPool(bufferSize int) Pool {
 	if bufferSize <= 0 {
 		bufferSize = DefaultBufferSize
 	}
-	p := &sync.Pool{New: func() interface{} {
-		return make([]byte, 0, bufferSize)
+	p := &pool{
+		bufSize: bufferSize,
+	}
+	p.p = &sync.Pool{New: func() interface{} {
+		return NewBuffer(p, make([]byte, bufferSize))
 	}}
-	return &pool{p: p, bufSize: bufferSize}
+	return p
 }
 
 func (p *pool) BufferSize() int {
 	return p.bufSize
 }
 
-func (p *pool) Put(bytes []byte) {
-	p.p.Put(bytes)
+func (p *pool) Put(buf *Buffer) {
+	p.p.Put(buf)
 }
 
-func (p *pool) Get(ctx context.Context) ([]byte, error) {
-	return p.p.Get().([]byte), nil
+func (p *pool) Get(ctx context.Context) (*Buffer, error) {
+	return p.p.Get().(*Buffer), nil
 }