Adapt buffer pool to recycle buffers only when safe

Author: Gilthoniel

pulsar/consumer_multitopic_test.go

@@ -332,7 +332,7 @@ func (dummyConnection) SendRequestNoWait(_ *pb.BaseCommand) error {
 	return nil
 }
 
-func (dummyConnection) WriteData(_ context.Context, _ internal.Buffer) {
+func (dummyConnection) WriteData(_ context.Context, _ *internal.SharedBuffer) {
 }
 
 func (dummyConnection) RegisterListener(_ uint64, _ internal.ConnectionListener) error {

pulsar/internal/batch_builder.go

@@ -30,7 +30,7 @@ import (
 )
 
 type BuffersPool interface {
-	GetBuffer() Buffer
+	GetBuffer() *SharedBuffer
 }
 
 // BatcherBuilderProvider defines func which returns the BatchBuilder.
@@ -71,7 +71,7 @@ type BatchBuilder interface {
 }
 
 type FlushBatch struct {
-	BatchData  Buffer
+	BatchData  *SharedBuffer
 	SequenceID uint64
 	Callbacks  []interface{}
 	Error      error
@@ -271,7 +271,7 @@ func (bc *batchContainer) Flush() *FlushBatch {
 
 	buffer := bc.buffersPool.GetBuffer()
 	if buffer == nil {
-		buffer = NewBuffer(int(uncompressedSize * 3 / 2))
+		buffer = NewSharedBuffer(NewBuffer(int(uncompressedSize * 3 / 2)))
 	}
 
 	sequenceID := uint64(0)

pulsar/internal/buffer.go

@@ -20,6 +20,7 @@ package internal
 import (
 	"encoding/binary"
 	"sync"
+	"sync/atomic"
 
 	log "github.com/sirupsen/logrus"
 )
@@ -227,14 +228,10 @@ func (b *buffer) Clear() {
 
 type BufferPool interface {
 	// Get returns a cleared buffer if any is available, otherwise nil.
-	Get() Buffer
+	Get() *SharedBuffer
 
 	// 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
+	Put(*SharedBuffer)
 }
 
 type synchronizedBufferPool struct {
@@ -247,23 +244,40 @@ func newBufferPool(pool *sync.Pool) synchronizedBufferPool {
 	return synchronizedBufferPool{pool: pool}
 }
 
-func (p synchronizedBufferPool) Get() Buffer {
+func (p synchronizedBufferPool) Get() *SharedBuffer {
 	buffer, ok := p.pool.Get().(Buffer)
 	if ok {
 		buffer.Clear()
+		return NewSharedBuffer(buffer)
+	}
+	return nil
+}
+
+func (p synchronizedBufferPool) Put(buffer *SharedBuffer) {
+	if !buffer.isCurrentlyShared.CompareAndSwap(true, false) {
+		p.pool.Put(buffer)
 	}
-	return buffer
 }
 
-func (p synchronizedBufferPool) Put(buffer Buffer) {
-	p.pool.Put(buffer)
+// SharedBuffer is a wrapper around a buffer that can keep track of a shared
+// reference of the buffer to later recycle it only when both references are
+// done with it.
+type SharedBuffer struct {
+	Buffer
+	isCurrentlyShared atomic.Bool
+}
+
+func NewSharedBuffer(buffer Buffer) *SharedBuffer {
+	return &SharedBuffer{Buffer: buffer}
 }
 
-func (p synchronizedBufferPool) Clone(b Buffer) Buffer {
-	newBuffer := p.Get()
-	if newBuffer == nil {
-		newBuffer = &buffer{}
+// Retain creates a new shared buffer backed by the same buffer but that will be
+// retained by the pool until the shared reference is also done.
+func (b *SharedBuffer) Retain() *SharedBuffer {
+	newBuffer := &SharedBuffer{
+		Buffer:            b.Buffer,
+		isCurrentlyShared: atomic.Bool{},
 	}
-	newBuffer.Write(b.ReadableSlice())
+	newBuffer.isCurrentlyShared.Store(true)
 	return newBuffer
 }

pulsar/internal/buffer_test.go

@@ -36,52 +36,58 @@ func TestBuffer(t *testing.T) {
 	assert.Equal(t, uint32(1024), b.Capacity())
 }
 
-func TestSynchronizedBufferPool_Clone_returnsNewlyAllocatedBuffer(t *testing.T) {
+func TestSynchronizedBufferPool_Get_returnsNilWhenEmpty(t *testing.T) {
 	pool := newBufferPool(&sync.Pool{})
+	assert.Nil(t, pool.Get())
+}
 
-	buffer := NewBuffer(1024)
-	buffer.Write([]byte{1, 2, 3})
+func TestSynchronizedBufferPool_Put_marksBufferAsNotSharedAnymore(t *testing.T) {
+	buffer := NewSharedBuffer(NewBuffer(1024)).Retain()
+	assert.True(t, buffer.isCurrentlyShared.Load())
 
-	res := pool.Clone(buffer)
-	assert.Equal(t, []byte{1, 2, 3}, res.ReadableSlice())
+	pool := newBufferPool(&sync.Pool{})
+	pool.Put(buffer)
+	assert.False(t, buffer.isCurrentlyShared.Load())
 }
 
-func TestSynchronizedBufferPool_Clone_returnsRecycledBuffer(t *testing.T) {
+func TestSynchronizedBufferPool_Put_recyclesSharedBuffer(t *testing.T) {
 	pool := newBufferPool(&sync.Pool{})
 
 	for range 100 {
-		buffer := NewBuffer(1024)
-		buffer.Write([]byte{1, 2, 3})
+		buffer := NewSharedBuffer(NewBuffer(1024)).Retain()
+		pool.Put(buffer)
 		pool.Put(buffer)
-	}
 
-	buffer := NewBuffer(1024)
-	buffer.Write([]byte{1, 2, 3})
+		if res := pool.Get(); res != nil {
+			return
+		}
+	}
 
-	res := pool.Clone(buffer)
-	assert.Equal(t, []byte{1, 2, 3}, res.ReadableSlice())
+	t.Fatal("pool is not recycling buffers")
 }
 
-// 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))
+func TestSynchronizedBufferPool_Put_recyclesBuffer(t *testing.T) {
+	pool := newBufferPool(&sync.Pool{})
+
+	for range 100 {
+		buffer := NewSharedBuffer(NewBuffer(1024))
+		pool.Put(buffer)
 
-	for range b.N {
-		newBuffer := pool.Clone(buffer)
-		pool.Put(newBuffer)
+		if res := pool.Get(); res != nil {
+			return
+		}
 	}
+
+	t.Fatal("pool is not recycling buffers")
 }
 
 // --- Helpers
 
 type capturingPool struct {
-	buffers []Buffer
+	buffers []*SharedBuffer
 }
 
-func (p *capturingPool) Get() Buffer {
+func (p *capturingPool) Get() *SharedBuffer {
 	if len(p.buffers) > 0 {
 		value := p.buffers[0]
 		p.buffers = p.buffers[1:]
@@ -90,10 +96,6 @@ func (p *capturingPool) Get() Buffer {
 	return nil
 }
 
-func (p *capturingPool) Put(value Buffer) {
+func (p *capturingPool) Put(value *SharedBuffer) {
 	p.buffers = append(p.buffers, value)
 }
-
-func (p *capturingPool) Clone(value Buffer) Buffer {
-	return value
-}

pulsar/internal/connection.go

@@ -79,7 +79,7 @@ type ConnectionListener interface {
 type Connection interface {
 	SendRequest(requestID uint64, req *pb.BaseCommand, callback func(*pb.BaseCommand, error))
 	SendRequestNoWait(req *pb.BaseCommand) error
-	WriteData(ctx context.Context, data Buffer)
+	WriteData(ctx context.Context, data *SharedBuffer)
 	RegisterListener(id uint64, listener ConnectionListener) error
 	UnregisterListener(id uint64)
 	AddConsumeHandler(id uint64, handler ConsumerHandler) error
@@ -132,7 +132,7 @@ type request struct {
 
 type dataRequest struct {
 	ctx  context.Context
-	data Buffer
+	data *SharedBuffer
 }
 
 type connection struct {
@@ -474,7 +474,7 @@ func (c *connection) runPingCheck(pingCheckTicker *time.Ticker) {
 	}
 }
 
-func (c *connection) WriteData(ctx context.Context, data Buffer) {
+func (c *connection) WriteData(ctx context.Context, data *SharedBuffer) {
 	written := false
 	defer func() {
 		if !written {

pulsar/internal/connection_test.go

@@ -38,7 +38,7 @@ func TestConnection_WriteData_recyclesBufferOnContextCanceled(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	cancel()
 
-	c.WriteData(ctx, NewBuffer(1024))
+	c.WriteData(ctx, NewSharedBuffer(NewBuffer(1024)))
 	assert.NotNil(t, pool.Get())
 }
 
@@ -49,7 +49,7 @@ func TestConnection_WriteData_recyclesBufferOnConnectionClosed(t *testing.T) {
 	c.state.Store(connectionClosed)
 	c.bufferPool = pool
 
-	c.WriteData(context.Background(), NewBuffer(1024))
+	c.WriteData(context.Background(), NewSharedBuffer(NewBuffer(1024)))
 	assert.NotNil(t, pool.Get())
 }
 
@@ -58,7 +58,7 @@ func TestConnection_WriteData_doNotRecyclesBufferWhenWritten(t *testing.T) {
 	c := makeTestConnection()
 	c.bufferPool = pool
 
-	c.WriteData(context.Background(), NewBuffer(1024))
+	c.WriteData(context.Background(), NewSharedBuffer(NewBuffer(1024)))
 	assert.Nil(t, pool.Get())
 }
 
@@ -67,7 +67,7 @@ func TestConnection_run_recyclesBufferOnceDone(t *testing.T) {
 	c := makeTestConnection()
 	c.bufferPool = pool
 
-	c.writeRequestsCh <- &dataRequest{ctx: context.Background(), data: NewBuffer(1024)}
+	c.writeRequestsCh <- &dataRequest{ctx: context.Background(), data: NewSharedBuffer(NewBuffer(1024))}
 	close(c.writeRequestsCh)
 
 	c.run()
@@ -82,7 +82,7 @@ func TestConnection_run_recyclesBufferEvenOnContextCanceled(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	cancel()
 
-	c.writeRequestsCh <- &dataRequest{ctx: ctx, data: NewBuffer(1024)}
+	c.writeRequestsCh <- &dataRequest{ctx: ctx, data: NewSharedBuffer(NewBuffer(1024))}
 	close(c.writeRequestsCh)
 
 	c.run()

pulsar/internal/key_based_batch_builder_test.go

@@ -33,7 +33,7 @@ import (
 type mockBufferPool struct {
 }
 
-func (m *mockBufferPool) GetBuffer() Buffer {
+func (m *mockBufferPool) GetBuffer() *SharedBuffer {
 	return nil
 }
 

pulsar/producer_partition.go

@@ -393,9 +393,12 @@ func (p *partitionProducer) grabCnx(assignedBrokerURL string) error {
 			// when resending pending batches, we update the sendAt timestamp to record the metric.
 			pi.Lock()
 			pi.sentAt = time.Now()
+			// Buffer is sent again to the connection to be processed in an
+			// asynchronous routine so we need to retain it again.
+			pi.buffer = pi.buffer.Retain()
 			pi.Unlock()
 			p.pendingQueue.Put(pi)
-			p._getConn().WriteData(pi.ctx, p.bufferPool.Clone(pi.buffer))
+			p._getConn().WriteData(pi.ctx, pi.buffer)
 
 			if pi == lastViewItem {
 				break
@@ -413,7 +416,7 @@ func (cc *connectionClosed) HasURL() bool {
 	return len(cc.assignedBrokerURL) > 0
 }
 
-func (p *partitionProducer) GetBuffer() internal.Buffer {
+func (p *partitionProducer) GetBuffer() *internal.SharedBuffer {
 	return p.bufferPool.Get()
 }
 
@@ -795,7 +798,7 @@ func (p *partitionProducer) internalSingleSend(
 
 	buffer := p.GetBuffer()
 	if buffer == nil {
-		buffer = internal.NewBuffer(int(payloadBuf.ReadableBytes() * 3 / 2))
+		buffer = internal.NewSharedBuffer(internal.NewBuffer(int(payloadBuf.ReadableBytes() * 3 / 2)))
 	}
 
 	sid := *mm.SequenceId
@@ -836,7 +839,7 @@ type pendingItem struct {
 	sync.Mutex
 	ctx           context.Context
 	cancel        context.CancelFunc
-	buffer        internal.Buffer
+	buffer        *internal.SharedBuffer
 	sequenceID    uint64
 	createdAt     time.Time
 	sentAt        time.Time
@@ -883,7 +886,7 @@ func (p *partitionProducer) internalFlushCurrentBatch() {
 	p.writeData(batchData, sequenceID, callbacks)
 }
 
-func (p *partitionProducer) writeData(buffer internal.Buffer, sequenceID uint64, callbacks []interface{}) {
+func (p *partitionProducer) writeData(buffer *internal.SharedBuffer, sequenceID uint64, callbacks []interface{}) {
 	select {
 	case <-p.ctx.Done():
 		for _, cb := range callbacks {
@@ -893,6 +896,11 @@ func (p *partitionProducer) writeData(buffer internal.Buffer, sequenceID uint64,
 		}
 		return
 	default:
+		// As the buffer will be sent to the connection in an asynchronous
+		// routine, we need to announce it to later recycle later only when we
+		// are done in both the connection and the producer.
+		buffer = buffer.Retain()
+
 		now := time.Now()
 		ctx, cancel := context.WithCancel(context.Background())
 		p.pendingQueue.Put(&pendingItem{
@@ -904,7 +912,8 @@ func (p *partitionProducer) writeData(buffer internal.Buffer, sequenceID uint64,
 			sequenceID:   sequenceID,
 			sendRequests: callbacks,
 		})
-		p._getConn().WriteData(ctx, p.bufferPool.Clone(buffer))
+
+		p._getConn().WriteData(ctx, buffer)
 	}
 }
 

pulsar/producer_test.go

@@ -2607,3 +2607,32 @@ func TestSelectConnectionForSameProducer(t *testing.T) {
 
 	client.Close()
 }
+
+func TestProducerPutsBufferBackInPool(t *testing.T) {
+	client, err := NewClient(ClientOptions{
+		URL: serviceURL,
+	})
+	assert.NoError(t, err)
+	defer client.Close()
+
+	testProducer, err := client.CreateProducer(ProducerOptions{Topic: newTopicName()})
+	assert.NoError(t, err)
+	defer testProducer.Close()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	for range 100 {
+		// Due to the sync.Pool not keeping all buffers, we need a few to make
+		// sure there is at least one available at the end
+		_, err = testProducer.Send(ctx, &ProducerMessage{Payload: []byte{42}})
+		assert.NoError(t, err)
+
+		if buffer := internal.GetDefaultBufferPool().Get(); buffer != nil {
+			// Ok!
+			return
+		}
+	}
+
+	t.Fatal("no buffer has been recycled")
+}