Fix sending buffer race by using proper reference counting (#1394)

Fixes #1371

### Motivation

The issue is that the buffer is sent to the connection when the data is flushed at this line: https://github.com/apache/pulsar-client-go/blob/0ab28c229e4dab2adb505a135325b6ede6e0e4f4/pulsar/producer_partition.go#L910

Because of this, both the data request in the connection and the pending item hold onto this buffer.

However, if the pending item fails due to a timeout or the producer closing, the buffer is returned to the buffer pool at this line: https://github.com/apache/pulsar-client-go/blob/0ab28c229e4dab2adb505a135325b6ede6e0e4f4/pulsar/producer_partition.go#L1747

Yet, the connection still has this buffer. When the connection tries to handle the data request, it will try to use a buffer that has already been returned. This causes a data race.

This fix implements proper reference counting to ensure buffers are only returned to the pool when all references are released, preventing the data race.

### Modifications

  - Added `Retain()` and `Release()` methods to the `Buffer` interface
  - Implemented atomic reference counting using `atomic.Int64`
  - Buffers are only returned to the pool when reference count reaches zero
  - Add new metric `pulsar_client_sending_buffers_count` to show the sending buffers count.
  - This PR also fixes serval minor issues like data/request cleanup when closing the connection

Author: RobertIndie

pulsar/consumer_multitopic_test.go

@@ -354,7 +354,7 @@ func (dummyConnection) ID() string {
 }
 
 func (dummyConnection) GetMaxMessageSize() int32 {
-	return 0
+	return 5 * 1024 * 1024 // 5MB
 }
 
 func (dummyConnection) Close() {

pulsar/internal/batch_builder.go

@@ -29,10 +29,6 @@ import (
 	"github.com/apache/pulsar-client-go/pulsar/log"
 )
 
-type BuffersPool interface {
-	GetBuffer() Buffer
-}
-
 // BatcherBuilderProvider defines func which returns the BatchBuilder.
 type BatcherBuilderProvider func(
 	maxMessages uint, maxBatchSize uint, maxMessageSize uint32, producerName string, producerID uint64,
@@ -269,10 +265,7 @@ func (bc *batchContainer) Flush() *FlushBatch {
 	uncompressedSize := bc.buffer.ReadableBytes()
 	bc.msgMetadata.UncompressedSize = &uncompressedSize
 
-	buffer := bc.buffersPool.GetBuffer()
-	if buffer == nil {
-		buffer = NewBuffer(int(uncompressedSize * 3 / 2))
-	}
+	buffer := bc.buffersPool.GetBuffer(int(uncompressedSize * 3 / 2))
 
 	sequenceID := uint64(0)
 	var err error

pulsar/internal/buffer.go

@@ -19,10 +19,17 @@ package internal
 
 import (
 	"encoding/binary"
+	"sync"
+	"sync/atomic"
 
 	log "github.com/sirupsen/logrus"
 )
 
+type BuffersPool interface {
+	GetBuffer(initSize int) Buffer
+	Put(buf Buffer)
+}
+
 // 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 {
@@ -69,15 +76,56 @@ type Buffer interface {
 	Resize(newSize uint32)
 	ResizeIfNeeded(spaceNeeded uint32)
 
+	Retain()
+	Release()
+	RefCnt() int64
+
 	// Clear will clear the current buffer data.
 	Clear()
 }
 
+type bufferPoolImpl struct {
+	sync.Pool
+}
+
+func NewBufferPool() BuffersPool {
+	return &bufferPoolImpl{
+		Pool: sync.Pool{},
+	}
+}
+
+func (p *bufferPoolImpl) GetBuffer(initSize int) Buffer {
+	sendingBuffersCount.Inc()
+	b, ok := p.Get().(*buffer)
+	if ok {
+		b.Clear()
+	} else {
+		b = &buffer{
+			data:      make([]byte, initSize),
+			readerIdx: 0,
+			writerIdx: 0,
+		}
+	}
+	b.pool = p
+	b.Retain()
+	return b
+}
+
+func (p *bufferPoolImpl) Put(buf Buffer) {
+	sendingBuffersCount.Dec()
+	if b, ok := buf.(*buffer); ok {
+		p.Pool.Put(b)
+	}
+}
+
 type buffer struct {
 	data []byte
 
 	readerIdx uint32
 	writerIdx uint32
+
+	refCnt atomic.Int64
+	pool   BuffersPool
 }
 
 // NewBuffer creates and initializes a new Buffer using buf as its initial contents.
@@ -213,7 +261,24 @@ func (b *buffer) Put(writerIdx uint32, s []byte) {
 	copy(b.data[writerIdx:], s)
 }
 
+func (b *buffer) Retain() {
+	b.refCnt.Add(1)
+}
+
+func (b *buffer) Release() {
+	if b.refCnt.Add(-1) == 0 {
+		if b.pool != nil {
+			b.pool.Put(b)
+		}
+	}
+}
+
+func (b *buffer) RefCnt() int64 {
+	return b.refCnt.Load()
+}
+
 func (b *buffer) Clear() {
 	b.readerIdx = 0
 	b.writerIdx = 0
+	b.refCnt.Store(0)
 }

pulsar/internal/connection.go

@@ -385,13 +385,20 @@ func (c *connection) failLeftRequestsWhenClose() {
 		// if other requests come in after the nil message
 		// then the RPC client will time out
 		ch <- nil
+		c.writeRequestsCh <- nil
 	}()
 	for req := range ch {
 		if nil == req {
 			break // we have drained the requests
 		}
 		c.internalSendRequest(req)
 	}
+	for req := range c.writeRequestsCh {
+		if nil == req {
+			break
+		}
+		req.data.Release()
+	}
 }
 
 func (c *connection) run() {
@@ -432,6 +439,7 @@ func (c *connection) run() {
 				return
 			}
 			c.internalWriteData(req.ctx, req.data)
+			req.data.Release()
 
 		case <-pingSendTicker.C:
 			c.sendPing()
@@ -457,13 +465,23 @@ func (c *connection) runPingCheck(pingCheckTicker *time.Ticker) {
 }
 
 func (c *connection) WriteData(ctx context.Context, data Buffer) {
+	writeToQueue := false
+	defer func() {
+		if !writeToQueue {
+			data.Release()
+		}
+	}()
 	select {
 	case c.writeRequestsCh <- &dataRequest{ctx: ctx, data: data}:
 		// Channel is not full
+		writeToQueue = true
 		return
 	case <-ctx.Done():
 		c.log.Debug("Write data context cancelled")
 		return
+	case <-c.closeCh:
+		c.log.Debug("Write data connection closed")
+		return
 	default:
 		// Channel full, fallback to probe if connection is closed
 	}
@@ -472,10 +490,14 @@ func (c *connection) WriteData(ctx context.Context, data Buffer) {
 		select {
 		case c.writeRequestsCh <- &dataRequest{ctx: ctx, data: data}:
 			// Successfully wrote on the channel
+			writeToQueue = true
 			return
 		case <-ctx.Done():
 			c.log.Debug("Write data context cancelled")
 			return
+		case <-c.closeCh:
+			c.log.Debug("Write data connection closed")
+			return
 		case <-time.After(100 * time.Millisecond):
 			// The channel is either:
 			// 1. blocked, in which case we need to wait until we have space
@@ -497,6 +519,8 @@ func (c *connection) internalWriteData(ctx context.Context, data Buffer) {
 	select {
 	case <-ctx.Done():
 		return
+	case <-c.closeCh:
+		return
 	default:
 		if _, err := c.cnx.Write(data.ReadableSlice()); err != nil {
 			c.log.WithError(err).Warn("Failed to write on connection")

pulsar/internal/key_based_batch_builder_test.go

@@ -30,13 +30,6 @@ import (
 	"google.golang.org/protobuf/proto"
 )
 
-type mockBufferPool struct {
-}
-
-func (m *mockBufferPool) GetBuffer() Buffer {
-	return nil
-}
-
 type mockEncryptor struct {
 }
 
@@ -53,7 +46,7 @@ func TestKeyBasedBatcherOrdering(t *testing.T) {
 		1,
 		pb.CompressionType_NONE,
 		compression.Level(0),
-		&mockBufferPool{},
+		&bufferPoolImpl{},
 		log.NewLoggerWithLogrus(logrus.StandardLogger()),
 		&mockEncryptor{},
 	)

pulsar/internal/metrics.go

@@ -21,6 +21,18 @@ import (
 	"github.com/prometheus/client_golang/prometheus"
 )
 
+var (
+	defaultConstLabels = map[string]string{
+		"client": "go",
+	}
+
+	sendingBuffersCount = prometheus.NewGauge(prometheus.GaugeOpts{
+		Name:        "pulsar_client_sending_buffers_count",
+		Help:        "Number of sending buffers",
+		ConstLabels: defaultConstLabels,
+	})
+)
+
 type Metrics struct {
 	metricsLevel      int
 	messagesPublished *prometheus.CounterVec
@@ -99,8 +111,9 @@ type LeveledMetrics struct {
 // NewMetricsProvider returns metrics registered to registerer.
 func NewMetricsProvider(metricsCardinality int, userDefinedLabels map[string]string,
 	registerer prometheus.Registerer) *Metrics {
-	constLabels := map[string]string{
-		"client": "go",
+	constLabels := make(map[string]string)
+	for k, v := range defaultConstLabels {
+		constLabels[k] = v
 	}
 	for k, v := range userDefinedLabels {
 		constLabels[k] = v
@@ -535,6 +548,7 @@ func NewMetricsProvider(metricsCardinality int, userDefinedLabels map[string]str
 			metrics.RPCRequestCount = are.ExistingCollector.(prometheus.Counter)
 		}
 	}
+	_ = registerer.Register(sendingBuffersCount)
 	return metrics
 }
 

pulsar/producer_partition.go

@@ -69,7 +69,7 @@ var (
 	ErrProducerBlockedQuotaExceeded = newError(ProducerBlockedQuotaExceededException, "producer blocked")
 	ErrProducerFenced               = newError(ProducerFenced, "producer fenced")
 
-	buffersPool     sync.Pool
+	buffersPool     internal.BuffersPool
 	sendRequestPool *sync.Pool
 )
 
@@ -86,6 +86,7 @@ func init() {
 			return &sendRequest{}
 		},
 	}
+	buffersPool = internal.NewBufferPool()
 }
 
 type partitionProducer struct {
@@ -96,7 +97,7 @@ type partitionProducer struct {
 	topic  string
 	log    log.Logger
 
-	conn         uAtomic.Value
+	conn         atomic.Pointer[internal.Connection]
 	cnxKeySuffix int32
 
 	options                  *ProducerOptions
@@ -363,7 +364,7 @@ func (p *partitionProducer) grabCnx(assignedBrokerURL string) error {
 		p.batchBuilder, err = provider(p.options.BatchingMaxMessages, p.options.BatchingMaxSize,
 			maxMessageSize, p.producerName, p.producerID, pb.CompressionType(p.options.CompressionType),
 			compression.Level(p.options.CompressionLevel),
-			p,
+			buffersPool,
 			p.log,
 			p.encryptor)
 		if err != nil {
@@ -393,6 +394,7 @@ func (p *partitionProducer) grabCnx(assignedBrokerURL string) error {
 			pi.Lock()
 			pi.sentAt = time.Now()
 			pi.Unlock()
+			pi.buffer.Retain() // Retain for writing to the connection
 			p.pendingQueue.Put(pi)
 			p._getConn().WriteData(pi.ctx, pi.buffer)
 
@@ -412,14 +414,6 @@ func (cc *connectionClosed) HasURL() bool {
 	return len(cc.assignedBrokerURL) > 0
 }
 
-func (p *partitionProducer) GetBuffer() internal.Buffer {
-	b, ok := buffersPool.Get().(internal.Buffer)
-	if ok {
-		b.Clear()
-	}
-	return b
-}
-
 func (p *partitionProducer) ConnectionClosed(closeProducer *pb.CommandCloseProducer) {
 	// Trigger reconnection in the produce goroutine
 	p.log.WithField("cnx", p._getConn().ID()).Warn("Connection was closed")
@@ -796,10 +790,7 @@ func (p *partitionProducer) internalSingleSend(
 	payloadBuf := internal.NewBuffer(len(compressedPayload))
 	payloadBuf.Write(compressedPayload)
 
-	buffer := p.GetBuffer()
-	if buffer == nil {
-		buffer = internal.NewBuffer(int(payloadBuf.ReadableBytes() * 3 / 2))
-	}
+	buffer := buffersPool.GetBuffer(int(payloadBuf.ReadableBytes() * 3 / 2))
 
 	sid := *mm.SequenceId
 	var useTxn bool
@@ -889,6 +880,7 @@ func (p *partitionProducer) internalFlushCurrentBatch() {
 func (p *partitionProducer) writeData(buffer internal.Buffer, sequenceID uint64, callbacks []interface{}) {
 	select {
 	case <-p.ctx.Done():
+		buffer.Release()
 		for _, cb := range callbacks {
 			if sr, ok := cb.(*sendRequest); ok {
 				sr.done(nil, ErrProducerClosed)
@@ -898,6 +890,7 @@ func (p *partitionProducer) writeData(buffer internal.Buffer, sequenceID uint64,
 	default:
 		now := time.Now()
 		ctx, cancel := context.WithCancel(context.Background())
+		buffer.Retain()
 		p.pendingQueue.Put(&pendingItem{
 			ctx:          ctx,
 			cancel:       cancel,
@@ -1744,7 +1737,7 @@ func (i *pendingItem) done(err error) {
 
 	i.isDone = true
 	// return the buffer to the pool after all callbacks have been called.
-	defer buffersPool.Put(i.buffer)
+	defer i.buffer.Release()
 	if i.flushCallback != nil {
 		i.flushCallback(err)
 	}
@@ -1757,16 +1750,16 @@ func (i *pendingItem) done(err error) {
 // _setConn sets the internal connection field of this partition producer atomically.
 // Note: should only be called by this partition producer when a new connection is available.
 func (p *partitionProducer) _setConn(conn internal.Connection) {
-	p.conn.Store(conn)
+	p.conn.Store(&conn)
 }
 
 // _getConn returns internal connection field of this partition producer atomically.
 // Note: should only be called by this partition producer before attempting to use the connection
 func (p *partitionProducer) _getConn() internal.Connection {
-	// Invariant: p.conn must be non-nil for the lifetime of the partitionProducer.
+	// Invariant: The conn must be non-nil for the lifetime of the partitionProducer.
 	//            For this reason we leave this cast unchecked and panic() if the
 	//            invariant is broken
-	return p.conn.Load().(internal.Connection)
+	return *p.conn.Load()
 }
 
 type chunkRecorder struct {