Recycle a buffer in several async routines

Author: Gilthoniel

pulsar/internal/buffer.go

@@ -19,10 +19,17 @@ package internal
 
 import (
 	"encoding/binary"
+	"sync"
 
 	log "github.com/sirupsen/logrus"
 )
 
+var defaultBufferPool = newBufferPool(&sync.Pool{})
+
+func GetDefaultBufferPool() synchronizedBufferPool {
+	return defaultBufferPool
+}
+
 // Buffer is a variable-sized buffer of bytes with Read and Write methods.
 // The zero value for Buffer is an empty buffer ready to use.
 type Buffer interface {
@@ -217,3 +224,46 @@ func (b *buffer) Clear() {
 	b.readerIdx = 0
 	b.writerIdx = 0
 }
+
+type BufferPool interface {
+	// Get returns a cleared buffer if any is available, otherwise nil.
+	Get() Buffer
+
+	// Put puts the buffer back to the pool and available for other routines.
+	Put(Buffer)
+
+	// Clone attempts to create a clone using a buffer from the pool, or returns
+	// a new one if necessary.
+	Clone(Buffer) Buffer
+}
+
+type synchronizedBufferPool struct {
+	pool *sync.Pool
+}
+
+var _ BufferPool = synchronizedBufferPool{}
+
+func newBufferPool(pool *sync.Pool) synchronizedBufferPool {
+	return synchronizedBufferPool{pool: pool}
+}
+
+func (p synchronizedBufferPool) Get() Buffer {
+	buffer, ok := p.pool.Get().(Buffer)
+	if ok {
+		buffer.Clear()
+	}
+	return buffer
+}
+
+func (p synchronizedBufferPool) Put(buffer Buffer) {
+	p.pool.Put(buffer)
+}
+
+func (p synchronizedBufferPool) Clone(b Buffer) Buffer {
+	newBuffer := p.Get()
+	if newBuffer == nil {
+		newBuffer = &buffer{}
+	}
+	newBuffer.Write(b.ReadableSlice())
+	return newBuffer
+}

pulsar/internal/buffer_test.go

@@ -18,6 +18,7 @@
 package internal
 
 import (
+	"sync"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -34,3 +35,65 @@ func TestBuffer(t *testing.T) {
 	assert.Equal(t, uint32(1019), b.WritableBytes())
 	assert.Equal(t, uint32(1024), b.Capacity())
 }
+
+func TestSynchronizedBufferPool_Clone_returnsNewlyAllocatedBuffer(t *testing.T) {
+	pool := newBufferPool(&sync.Pool{})
+
+	buffer := NewBuffer(1024)
+	buffer.Write([]byte{1, 2, 3})
+
+	res := pool.Clone(buffer)
+	assert.Equal(t, []byte{1, 2, 3}, res.ReadableSlice())
+}
+
+func TestSynchronizedBufferPool_Clone_returnsRecycledBuffer(t *testing.T) {
+	pool := newBufferPool(&sync.Pool{})
+
+	for range 100 {
+		buffer := NewBuffer(1024)
+		buffer.Write([]byte{1, 2, 3})
+		pool.Put(buffer)
+	}
+
+	buffer := NewBuffer(1024)
+	buffer.Write([]byte{1, 2, 3})
+
+	res := pool.Clone(buffer)
+	assert.Equal(t, []byte{1, 2, 3}, res.ReadableSlice())
+}
+
+// BenchmarkBufferPool_Clone demonstrates the cloning of a buffer without
+// allocation if the pool is filled making the process very efficient.
+func BenchmarkBufferPool_Clone(b *testing.B) {
+	pool := GetDefaultBufferPool()
+	buffer := NewBuffer(1024)
+	buffer.Write(make([]byte, 1024))
+
+	for range b.N {
+		newBuffer := pool.Clone(buffer)
+		pool.Put(newBuffer)
+	}
+}
+
+// --- Helpers
+
+type capturingPool struct {
+	buffers []Buffer
+}
+
+func (p *capturingPool) Get() Buffer {
+	if len(p.buffers) > 0 {
+		value := p.buffers[0]
+		p.buffers = p.buffers[1:]
+		return value
+	}
+	return nil
+}
+
+func (p *capturingPool) Put(value Buffer) {
+	p.buffers = append(p.buffers, value)
+}
+
+func (p *capturingPool) Clone(value Buffer) Buffer {
+	return value
+}

pulsar/internal/connection.go

@@ -150,6 +150,7 @@ type connection struct {
 	logicalAddr  *url.URL
 	physicalAddr *url.URL
 
+	bufferPool      BufferPool
 	writeBufferLock sync.Mutex
 	writeBuffer     Buffer
 	reader          *connectionReader
@@ -199,6 +200,7 @@ func newConnection(opts connectionOptions) *connection {
 		keepAliveInterval:    opts.keepAliveInterval,
 		logicalAddr:          opts.logicalAddr,
 		physicalAddr:         opts.physicalAddr,
+		bufferPool:           GetDefaultBufferPool(),
 		writeBuffer:          NewBuffer(4096),
 		log:                  opts.logger.SubLogger(log.Fields{"remote_addr": opts.physicalAddr}),
 		pendingReqs:          make(map[uint64]*request),
@@ -433,6 +435,9 @@ func (c *connection) run() {
 			}
 			c.internalWriteData(req.ctx, req.data)
 
+			// Write request is fully processed so we can release the buffer.
+			c.bufferPool.Put(req.data)
+
 		case <-pingSendTicker.C:
 			c.sendPing()
 		}
@@ -457,9 +462,18 @@ func (c *connection) runPingCheck(pingCheckTicker *time.Ticker) {
 }
 
 func (c *connection) WriteData(ctx context.Context, data Buffer) {
+	written := false
+	defer func() {
+		if !written {
+			// Buffer has failed to be written and can now be reused.
+			c.bufferPool.Put(data)
+		}
+	}()
+
 	select {
 	case c.writeRequestsCh <- &dataRequest{ctx: ctx, data: data}:
 		// Channel is not full
+		written = true
 		return
 	case <-ctx.Done():
 		c.log.Debug("Write data context cancelled")
@@ -472,6 +486,7 @@ func (c *connection) WriteData(ctx context.Context, data Buffer) {
 		select {
 		case c.writeRequestsCh <- &dataRequest{ctx: ctx, data: data}:
 			// Successfully wrote on the channel
+			written = true
 			return
 		case <-ctx.Done():
 			c.log.Debug("Write data context cancelled")

pulsar/internal/connection_test.go

@@ -0,0 +1,141 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package internal
+
+import (
+	"context"
+	"io"
+	"net"
+	"testing"
+	"time"
+
+	"github.com/apache/pulsar-client-go/pulsar/log"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestConnection_WriteData_recyclesBufferOnContextCanceled(t *testing.T) {
+	pool := &capturingPool{}
+	c := makeTestConnection()
+	c.writeRequestsCh = make(chan *dataRequest)
+	c.bufferPool = pool
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	c.WriteData(ctx, NewBuffer(1024))
+	assert.NotNil(t, pool.Get())
+}
+
+func TestConnection_WriteData_recyclesBufferOnConnectionClosed(t *testing.T) {
+	pool := &capturingPool{}
+	c := makeTestConnection()
+	c.writeRequestsCh = make(chan *dataRequest)
+	c.state.Store(connectionClosed)
+	c.bufferPool = pool
+
+	c.WriteData(context.Background(), NewBuffer(1024))
+	assert.NotNil(t, pool.Get())
+}
+
+func TestConnection_WriteData_doNotRecyclesBufferWhenWritten(t *testing.T) {
+	pool := &capturingPool{}
+	c := makeTestConnection()
+	c.bufferPool = pool
+
+	c.WriteData(context.Background(), NewBuffer(1024))
+	assert.Nil(t, pool.Get())
+}
+
+func TestConnection_run_recyclesBufferOnceDone(t *testing.T) {
+	pool := &capturingPool{}
+	c := makeTestConnection()
+	c.bufferPool = pool
+
+	c.writeRequestsCh <- &dataRequest{ctx: context.Background(), data: NewBuffer(1024)}
+	close(c.writeRequestsCh)
+
+	c.run()
+	assert.NotNil(t, pool.Get())
+}
+
+func TestConnection_run_recyclesBufferEvenOnContextCanceled(t *testing.T) {
+	pool := &capturingPool{}
+	c := makeTestConnection()
+	c.bufferPool = pool
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	c.writeRequestsCh <- &dataRequest{ctx: ctx, data: NewBuffer(1024)}
+	close(c.writeRequestsCh)
+
+	c.run()
+	assert.NotNil(t, pool.Get())
+}
+
+// --- Helpers
+
+func makeTestConnection() *connection {
+	c := &connection{
+		log:               log.DefaultNopLogger(),
+		keepAliveInterval: 30 * time.Second,
+		writeRequestsCh:   make(chan *dataRequest, 10),
+		closeCh:           make(chan struct{}),
+		cnx:               happyCnx{},
+		metrics:           NewMetricsProvider(0, make(map[string]string), prometheus.DefaultRegisterer),
+		bufferPool:        GetDefaultBufferPool(),
+	}
+	c.reader = newConnectionReader(c)
+	c.state.Store(connectionReady)
+	return c
+}
+
+type happyCnx struct{}
+
+func (happyCnx) Read(b []byte) (n int, err error) {
+	return 0, io.EOF
+}
+
+func (happyCnx) Write(b []byte) (n int, err error) {
+	return len(b), nil
+}
+
+func (happyCnx) Close() error {
+	return nil
+}
+
+func (happyCnx) LocalAddr() net.Addr {
+	return &net.IPAddr{IP: net.ParseIP("127.0.0.1")}
+}
+
+func (happyCnx) RemoteAddr() net.Addr {
+	return &net.IPAddr{IP: net.ParseIP("127.0.0.1")}
+}
+
+func (happyCnx) SetDeadline(t time.Time) error {
+	return nil
+}
+
+func (happyCnx) SetReadDeadline(t time.Time) error {
+	return nil
+}
+
+func (happyCnx) SetWriteDeadline(t time.Time) error {
+	return nil
+}