fix: potential data race

Author: gunli

pulsar/consumer_multitopic_test.go

@@ -18,18 +18,20 @@
 package pulsar
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"strings"
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/assert"
+
 	"github.com/apache/pulsar-client-go/pulsar/internal"
 	pb "github.com/apache/pulsar-client-go/pulsar/internal/pulsar_proto"
 	"github.com/apache/pulsar-client-go/pulsaradmin"
 	"github.com/apache/pulsar-client-go/pulsaradmin/pkg/admin/config"
 	"github.com/apache/pulsar-client-go/pulsaradmin/pkg/utils"
-	"github.com/stretchr/testify/assert"
 )
 
 func TestMultiTopicConsumerReceive(t *testing.T) {
@@ -330,7 +332,7 @@ func (dummyConnection) SendRequestNoWait(_ *pb.BaseCommand) error {
 	return nil
 }
 
-func (dummyConnection) WriteData(_ internal.Buffer) {
+func (dummyConnection) WriteData(_ context.Context, _ internal.Buffer) {
 }
 
 func (dummyConnection) RegisterListener(_ uint64, _ internal.ConnectionListener) error {

pulsar/internal/connection.go

@@ -18,6 +18,7 @@
 package internal
 
 import (
+	"context"
 	"crypto/tls"
 	"crypto/x509"
 	"errors"
@@ -78,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(data Buffer)
+	WriteData(ctx context.Context, data Buffer)
 	RegisterListener(id uint64, listener ConnectionListener) error
 	UnregisterListener(id uint64)
 	AddConsumeHandler(id uint64, handler ConsumerHandler) error
@@ -450,12 +451,14 @@ func (c *connection) runPingCheck(pingCheckTicker *time.Ticker) {
 	}
 }
 
-func (c *connection) WriteData(data Buffer) {
+func (c *connection) WriteData(ctx context.Context, data Buffer) {
 	select {
 	case c.writeRequestsCh <- data:
 		// Channel is not full
 		return
-
+	case <-ctx.Done():
+		c.log.Debug("Write data context cancelled")
+		return
 	default:
 		// Channel full, fallback to probe if connection is closed
 	}
@@ -465,7 +468,9 @@ func (c *connection) WriteData(data Buffer) {
 		case c.writeRequestsCh <- data:
 			// Successfully wrote on the channel
 			return
-
+		case <-ctx.Done():
+			c.log.Debug("Write data context cancelled")
+			return
 		case <-time.After(100 * time.Millisecond):
 			// The channel is either:
 			// 1. blocked, in which case we need to wait until we have space

pulsar/producer_partition.go

@@ -394,7 +394,7 @@ func (p *partitionProducer) grabCnx(assignedBrokerURL string) error {
 			pi.sentAt = time.Now()
 			pi.Unlock()
 			p.pendingQueue.Put(pi)
-			p._getConn().WriteData(pi.buffer)
+			p._getConn().WriteData(pi.ctx, pi.buffer)
 
 			if pi == lastViewItem {
 				break
@@ -837,6 +837,8 @@ func (p *partitionProducer) internalSingleSend(
 
 type pendingItem struct {
 	sync.Mutex
+	ctx           context.Context
+	cancel        context.CancelFunc
 	buffer        internal.Buffer
 	sequenceID    uint64
 	createdAt     time.Time
@@ -846,6 +848,21 @@ type pendingItem struct {
 	flushCallback func(err error)
 }
 
+func (i *pendingItem) done(err error) {
+	if i.isDone {
+		return
+	}
+	i.isDone = true
+	buffersPool.Put(i.buffer)
+	if i.flushCallback != nil {
+		i.flushCallback(err)
+	}
+
+	if i.cancel != nil {
+		i.cancel()
+	}
+}
+
 func (p *partitionProducer) internalFlushCurrentBatch() {
 	if p.batchBuilder == nil {
 		// batch is not enabled
@@ -895,14 +912,17 @@ func (p *partitionProducer) writeData(buffer internal.Buffer, sequenceID uint64,
 		return
 	default:
 		now := time.Now()
+		ctx, cancel := context.WithCancel(context.Background())
 		p.pendingQueue.Put(&pendingItem{
+			ctx:          ctx,
+			cancel:       cancel,
 			createdAt:    now,
 			sentAt:       now,
 			buffer:       buffer,
 			sequenceID:   sequenceID,
 			sendRequests: callbacks,
 		})
-		p._getConn().WriteData(buffer)
+		p._getConn().WriteData(ctx, buffer)
 	}
 }
 
@@ -1579,14 +1599,14 @@ type sendRequest struct {
 	uuid             string
 	chunkRecorder    *chunkRecorder
 
-	/// resource management
+	// resource management
 
 	memLimit          internal.MemoryLimitController
 	reservedMem       int64
 	semaphore         internal.Semaphore
 	reservedSemaphore int
 
-	/// convey settable state
+	// convey settable state
 
 	sendAsBatch         bool
 	transaction         *transaction
@@ -1659,7 +1679,7 @@ func (sr *sendRequest) done(msgID MessageID, err error) {
 }
 
 func (p *partitionProducer) blockIfQueueFull() bool {
-	//DisableBlockIfQueueFull == false means enable block
+	// DisableBlockIfQueueFull == false means enable block
 	return !p.options.DisableBlockIfQueueFull
 }
 
@@ -1732,17 +1752,6 @@ type flushRequest struct {
 	err    error
 }
 
-func (i *pendingItem) done(err error) {
-	if i.isDone {
-		return
-	}
-	i.isDone = true
-	buffersPool.Put(i.buffer)
-	if i.flushCallback != nil {
-		i.flushCallback(err)
-	}
-}
-
 // _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) {