[Producer] respect context cancellation in Flush

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++ {

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() {}

pulsar/producer.go

@@ -237,10 +237,11 @@ 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
 
+	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.

pulsar/producer_impl.go

@@ -346,6 +346,19 @@ func (p *producer) Flush() error {
 	return nil
 }
 
+func (p *producer) FlushWithCtx(ctx context.Context) error {
+	p.RLock()
+	defer p.RUnlock()
+
+	for _, pp := range p.producers {
+		if err := pp.FlushWithCtx(ctx); err != nil {
+			return err
+		}
+
+	}
+	return nil
+}
+
 func (p *producer) Close() {
 	p.closeOnce.Do(func() {
 		p.stopDiscovery()

pulsar/producer_partition.go

@@ -1433,6 +1433,26 @@ func (p *partitionProducer) Flush() error {
 	return flushReq.err
 }
 
+func (p *partitionProducer) FlushWithCtx(ctx context.Context) error {
+	flushReq := &flushRequest{
+		doneCh: make(chan struct{}),
+		err:    nil,
+	}
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case p.cmdChan <- flushReq:
+	}
+
+	// wait for the flush request to complete
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-flushReq.doneCh:
+		return flushReq.err
+	}
+}
+
 func (p *partitionProducer) getProducerState() producerState {
 	return producerState(p.state.Load())
 }

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())
 

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++ {