Repository: apache/pulsar-client-go
Pull request #1165 (2024-01-31 19:11:06)

[Producer] respect context cancellation in Flush
### Motivation

The producer's `Flush` method does not respect context cancellation. If the caller's context get's cancelled, it will have to wait for the producer to finish flushing.

### Modifications

This change adds a `FlushWithCtx` method which takes a context and selects on two channels.

### Verifying this change

- [x] Make sure that the change passes the CI checks.

This change is already covered by existing tests. See code changes for relevant tests.

### Does this pull request potentially affect one of the following parts:

*If `yes` was chosen, please highlight the changes*

  - Dependencies (does it add or upgrade a dependency): no
  - The public API: no
  - The schema: no
  - The default values of configurations: no
  - The wire protocol: no

### Documentation

  - Does this pull request introduce a new feature? no
  - If yes, how is the feature documented? not applicable
  - If a feature is not applicable for documentation, explain why? It's a very trivial change. It adds a common pattern in Go.
  - If a feature is not documented yet in this PR, please create a followup issue for adding the documentation


jayshrivastava: @RobertIndie Would you mind taking a look at this please?

Patch:
diff --git a/pulsar/internal/compression/zstd_cgo.go b/pulsar/internal/compression/zstd_cgo.go
index 25429e2..dde54ae 100644
--- a/pulsar/internal/compression/zstd_cgo.go
+++ b/pulsar/internal/compression/zstd_cgo.go
@@ -25,19 +25,23 @@
 package compression
 
 import (
+	"sync"
+
 	"github.com/DataDog/zstd"
 	log "github.com/sirupsen/logrus"
 )
 
 type zstdCGoProvider struct {
-	ctx       zstd.Ctx
+	ctxPool   sync.Pool
 	level     Level
 	zstdLevel int
 }
 
 func newCGoZStdProvider(level Level) Provider {
 	z := &zstdCGoProvider{
-		ctx: zstd.NewCtx(),
+		ctxPool: sync.Pool{New: func() any {
+			return zstd.NewCtx()
+		}},
 	}
 
 	switch level {
@@ -61,7 +65,9 @@ func (z *zstdCGoProvider) CompressMaxSize(originalSize int) int {
 }
 
 func (z *zstdCGoProvider) Compress(dst, src []byte) []byte {
-	out, err := z.ctx.CompressLevel(dst, src, z.zstdLevel)
+	ctx := z.ctxPool.Get().(zstd.Ctx)
+	defer z.ctxPool.Put(ctx)
+	out, err := ctx.CompressLevel(dst, src, z.zstdLevel)
 	if err != nil {
 		log.WithError(err).Fatal("Failed to compress")
 	}
@@ -70,7 +76,9 @@ func (z *zstdCGoProvider) Compress(dst, src []byte) []byte {
 }
 
 func (z *zstdCGoProvider) Decompress(dst, src []byte, originalSize int) ([]byte, error) {
-	return z.ctx.Decompress(dst, src)
+	ctx := z.ctxPool.Get().(zstd.Ctx)
+	defer z.ctxPool.Put(ctx)
+	return ctx.Decompress(dst, src)
 }
 
 func (z *zstdCGoProvider) Close() error {
diff --git a/pulsar/internal/utils.go b/pulsar/internal/utils.go
index 9378d9d..2dc8210 100644
--- a/pulsar/internal/utils.go
+++ b/pulsar/internal/utils.go
@@ -40,7 +40,7 @@ func TimestampMillis(t time.Time) uint64 {
 // GetAndAdd perform atomic read and update
 func GetAndAdd(n *uint64, diff uint64) uint64 {
 	for {
-		v := *n
+		v := atomic.LoadUint64(n)
 		if atomic.CompareAndSwapUint64(n, v, v+diff) {
 			return v
 		}
diff --git a/pulsar/producer.go b/pulsar/producer.go
index 70d152c..f8013a1 100644
--- a/pulsar/producer.go
+++ b/pulsar/producer.go
@@ -237,10 +237,13 @@ type Producer interface {
 	// return the last sequence id published by this producer.
 	LastSequenceID() int64
 
-	// Flush all the messages buffered in the client and wait until all messages have been successfully
-	// persisted.
+	// Deprecated: Use `FlushWithCtx()` instead.
 	Flush() error
 
+	// Flush all the messages buffered in the client and wait until all messageshave been successfully
+	// persisted.
+	FlushWithCtx(ctx context.Context) error
+
 	// Close the producer and releases resources allocated
 	// No more writes will be accepted from this producer. Waits until all pending write request are persisted. In case
 	// of errors, pending writes will not be retried.
diff --git a/pulsar/producer_impl.go b/pulsar/producer_impl.go
index 3c45b59..ca92310 100644
--- a/pulsar/producer_impl.go
+++ b/pulsar/producer_impl.go
@@ -334,11 +334,15 @@ func (p *producer) LastSequenceID() int64 {
 }
 
 func (p *producer) Flush() error {
+	return p.FlushWithCtx(context.Background())
+}
+
+func (p *producer) FlushWithCtx(ctx context.Context) error {
 	p.RLock()
 	defer p.RUnlock()
 
 	for _, pp := range p.producers {
-		if err := pp.Flush(); err != nil {
+		if err := pp.FlushWithCtx(ctx); err != nil {
 			return err
 		}
 
diff --git a/pulsar/producer_partition.go b/pulsar/producer_partition.go
index 1b79053..fbcc5b9 100755
--- a/pulsar/producer_partition.go
+++ b/pulsar/producer_partition.go
@@ -1422,15 +1422,27 @@ func (p *partitionProducer) LastSequenceID() int64 {
 }
 
 func (p *partitionProducer) Flush() error {
+	return p.FlushWithCtx(context.Background())
+}
+
+func (p *partitionProducer) FlushWithCtx(ctx context.Context) error {
 	flushReq := &flushRequest{
 		doneCh: make(chan struct{}),
 		err:    nil,
 	}
-	p.cmdChan <- flushReq
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case p.cmdChan <- flushReq:
+	}
 
 	// wait for the flush request to complete
-	<-flushReq.doneCh
-	return flushReq.err
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-flushReq.doneCh:
+		return flushReq.err
+	}
 }
 
 func (p *partitionProducer) getProducerState() producerState {


Test patch:
diff --git a/pulsar/consumer_test.go b/pulsar/consumer_test.go
index df70b0d..d66e237 100644
--- a/pulsar/consumer_test.go
+++ b/pulsar/consumer_test.go
@@ -983,7 +983,7 @@ func TestConsumerBatchCumulativeAck(t *testing.T) {
 	}
 	wg.Wait()
 
-	err = producer.Flush()
+	err = producer.FlushWithCtx(context.Background())
 	assert.NoError(t, err)
 
 	// send another batch
@@ -1218,7 +1218,7 @@ func TestConsumerCompressionWithBatches(t *testing.T) {
 		}, nil)
 	}
 
-	producer.Flush()
+	producer.FlushWithCtx(context.Background())
 
 	for i := 0; i < N; i++ {
 		msg, err := consumer.Receive(ctx)
@@ -3932,7 +3932,7 @@ func runBatchIndexAckTest(t *testing.T, ackWithResponse bool, cumulative bool, o
 			log.Printf("Sent to %v:%d:%d", id, id.BatchIdx(), id.BatchSize())
 		})
 	}
-	assert.Nil(t, producer.Flush())
+	assert.Nil(t, producer.FlushWithCtx(context.Background()))
 
 	msgIds := make([]MessageID, BatchingMaxSize)
 	for i := 0; i < BatchingMaxSize; i++ {
diff --git a/pulsar/internal/pulsartracing/producer_interceptor_test.go b/pulsar/internal/pulsartracing/producer_interceptor_test.go
index 8d8e696..1c2c712 100644
--- a/pulsar/internal/pulsartracing/producer_interceptor_test.go
+++ b/pulsar/internal/pulsartracing/producer_interceptor_test.go
@@ -67,4 +67,8 @@ func (p *mockProducer) Flush() error {
 	return nil
 }
 
+func (p *mockProducer) FlushWithCtx(ctx context.Context) error {
+	return nil
+}
+
 func (p *mockProducer) Close() {}
diff --git a/pulsar/producer_test.go b/pulsar/producer_test.go
index 0d74cde..ba59115 100644
--- a/pulsar/producer_test.go
+++ b/pulsar/producer_test.go
@@ -159,7 +159,7 @@ func TestProducerAsyncSend(t *testing.T) {
 		assert.NoError(t, err)
 	}
 
-	err = producer.Flush()
+	err = producer.FlushWithCtx(context.Background())
 	assert.Nil(t, err)
 
 	wg.Wait()
@@ -220,7 +220,7 @@ func TestProducerFlushDisableBatching(t *testing.T) {
 		assert.NoError(t, err)
 	}
 
-	err = producer.Flush()
+	err = producer.FlushWithCtx(context.Background())
 	assert.Nil(t, err)
 
 	wg.Wait()
@@ -387,7 +387,7 @@ func TestFlushInProducer(t *testing.T) {
 		})
 		assert.Nil(t, err)
 	}
-	err = producer.Flush()
+	err = producer.FlushWithCtx(context.Background())
 	assert.Nil(t, err)
 	wg.Wait()
 
@@ -429,7 +429,7 @@ func TestFlushInProducer(t *testing.T) {
 		assert.Nil(t, err)
 	}
 
-	err = producer.Flush()
+	err = producer.FlushWithCtx(context.Background())
 	assert.Nil(t, err)
 	wg.Wait()
 
@@ -500,7 +500,7 @@ func TestFlushInPartitionedProducer(t *testing.T) {
 	}
 
 	// After flush, should be able to consume.
-	err = producer.Flush()
+	err = producer.FlushWithCtx(context.Background())
 	assert.Nil(t, err)
 
 	wg.Wait()
@@ -1717,7 +1717,7 @@ func TestMultipleSchemaOfKeyBasedBatchProducerConsumer(t *testing.T) {
 		}
 
 	}
-	producer.Flush()
+	producer.FlushWithCtx(context.Background())
 
 	//// create consumer
 	consumer, err := client.Subscribe(ConsumerOptions{
@@ -1808,7 +1808,7 @@ func TestMultipleSchemaProducerConsumer(t *testing.T) {
 			assert.NotNil(t, id)
 		})
 	}
-	producer.Flush()
+	producer.FlushWithCtx(context.Background())
 
 	//// create consumer
 	consumer, err := client.Subscribe(ConsumerOptions{
@@ -2027,9 +2027,9 @@ func TestMemLimitRejectProducerMessages(t *testing.T) {
 	assert.ErrorContains(t, err, getResultStr(ClientMemoryBufferIsFull))
 
 	// flush pending msg
-	err = producer1.Flush()
+	err = producer1.FlushWithCtx(context.Background())
 	assert.NoError(t, err)
-	err = producer2.Flush()
+	err = producer2.FlushWithCtx(context.Background())
 	assert.NoError(t, err)
 	assert.Equal(t, int64(0), c.(*client).memLimit.CurrentUsage())
 
@@ -2118,9 +2118,9 @@ func TestMemLimitRejectProducerMessagesWithSchema(t *testing.T) {
 	assert.ErrorContains(t, err, getResultStr(ClientMemoryBufferIsFull))
 
 	// flush pending msg
-	err = producer1.Flush()
+	err = producer1.FlushWithCtx(context.Background())
 	assert.NoError(t, err)
-	err = producer2.Flush()
+	err = producer2.FlushWithCtx(context.Background())
 	assert.NoError(t, err)
 	assert.Equal(t, int64(0), c.(*client).memLimit.CurrentUsage())
 
@@ -2244,7 +2244,7 @@ func TestMemLimitContextCancel(t *testing.T) {
 	cancel()
 	wg.Wait()
 
-	err = producer.Flush()
+	err = producer.FlushWithCtx(context.Background())
 	assert.NoError(t, err)
 	assert.Equal(t, int64(0), c.(*client).memLimit.CurrentUsage())
 
@@ -2357,6 +2357,34 @@ func TestFailPendingMessageWithClose(t *testing.T) {
 	assert.Equal(t, 0, partitionProducerImp.pendingQueue.Size())
 }
 
+func TestSendConcurrently(t *testing.T) {
+	client, err := NewClient(ClientOptions{
+		URL: lookupURL,
+	})
+	assert.NoError(t, err)
+	defer client.Close()
+	testProducer, err := client.CreateProducer(ProducerOptions{
+		Topic:            newTopicName(),
+		CompressionType:  ZSTD,
+		CompressionLevel: Better,
+		DisableBatching:  true,
+	})
+	assert.NoError(t, err)
+
+	var wg sync.WaitGroup
+	for i := 0; i < 100; i++ {
+		wg.Add(1)
+		go func() {
+			_, err := testProducer.Send(context.Background(), &ProducerMessage{
+				Payload: make([]byte, 100),
+			})
+			assert.NoError(t, err)
+			wg.Done()
+		}()
+	}
+	wg.Wait()
+}
+
 type pendingQueueWrapper struct {
 	pendingQueue   internal.BlockingQueue
 	writtenBuffers *[]internal.Buffer
diff --git a/pulsar/reader_test.go b/pulsar/reader_test.go
index ec10f8f..c8228a7 100644
--- a/pulsar/reader_test.go
+++ b/pulsar/reader_test.go
@@ -277,7 +277,7 @@ func TestReaderOnSpecificMessageWithBatching(t *testing.T) {
 		})
 	}
 
-	err = producer.Flush()
+	err = producer.FlushWithCtx(context.Background())
 	assert.NoError(t, err)
 
 	// create reader on 5th message (not included)
@@ -353,7 +353,7 @@ func TestReaderOnLatestWithBatching(t *testing.T) {
 		})
 	}
 
-	err = producer.Flush()
+	err = producer.FlushWithCtx(context.Background())
 	assert.NoError(t, err)
 
 	// create reader on 5th message (not included)
@@ -592,7 +592,7 @@ func TestReaderSeek(t *testing.T) {
 			seekID = id
 		}
 	}
-	err = producer.Flush()
+	err = producer.FlushWithCtx(context.Background())
 	assert.NoError(t, err)
 
 	for i := 0; i < N; i++ {