contrib/confluentinc/confluent-kafka-go: fix goroutine leak in Produce

contrib/confluentinc/confluent-kafka-go/kafka.v2/kafka.go

@@ -334,31 +334,43 @@ func (p *Producer) Close() {
 func (p *Producer) Produce(msg *kafka.Message, deliveryChan chan kafka.Event) error {
 	span := p.startSpan(msg)
 
+	var errChan chan error
+
 	// if the user has selected a delivery channel, we will wrap it and
-	// wait for the delivery event to finish the span
+	// wait for the delivery event to finish the span.
+	// in case the Produce call returns an error, we won't receive anything
+	// in the deliveryChan, so we use errChan to make sure this goroutine exits.
 	if deliveryChan != nil {
+		errChan = make(chan error, 1)
 		oldDeliveryChan := deliveryChan
 		deliveryChan = make(chan kafka.Event)
 		go func() {
 			var err error
-			evt := <-deliveryChan
-			if msg, ok := evt.(*kafka.Message); ok {
-				// delivery errors are returned via TopicPartition.Error
-				err = msg.TopicPartition.Error
-				trackProduceOffsets(p.cfg.dataStreamsEnabled, msg, err)
+			select {
+			case evt := <-deliveryChan:
+				if msg, ok := evt.(*kafka.Message); ok {
+					// delivery errors are returned via TopicPartition.Error
+					err = msg.TopicPartition.Error
+					trackProduceOffsets(p.cfg.dataStreamsEnabled, msg, err)
+				}
+				oldDeliveryChan <- evt
+
+			case e := <-errChan:
+				err = e
 			}
 			span.Finish(tracer.WithError(err))
-			oldDeliveryChan <- evt
 		}()
 	}
 
 	setProduceCheckpoint(p.cfg.dataStreamsEnabled, p.libraryVersion, msg)
 	err := p.Producer.Produce(msg, deliveryChan)
-	// with no delivery channel or enqueue error, finish immediately
-	if err != nil || deliveryChan == nil {
-		span.Finish(tracer.WithError(err))
+	if err != nil {
+		if deliveryChan != nil {
+			errChan <- err
+		} else {
+			span.Finish(tracer.WithError(err))
+		}
 	}
-
 	return err
 }
 

contrib/confluentinc/confluent-kafka-go/kafka.v2/kafka_test.go

@@ -23,6 +23,7 @@ import (
 	"github.com/confluentinc/confluent-kafka-go/v2/kafka"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"go.uber.org/goleak"
 )
 
 var (
@@ -394,3 +395,41 @@ func TestNamingSchema(t *testing.T) {
 	}
 	namingschematest.NewKafkaTest(genSpans)(t)
 }
+
+// Test we don't leak goroutines and properly close the span when Produce returns an error.
+func TestProduceError(t *testing.T) {
+	defer goleak.VerifyNone(t)
+
+	mt := mocktracer.Start()
+	defer mt.Stop()
+
+	// first write a message to the topic
+	p, err := NewProducer(&kafka.ConfigMap{
+		"bootstrap.servers":   "127.0.0.1:9092",
+		"go.delivery.reports": true,
+	})
+	require.NoError(t, err)
+	defer p.Close()
+
+	// this empty message should cause an error in the Produce call.
+	topic := ""
+	msg := &kafka.Message{
+		TopicPartition: kafka.TopicPartition{
+			Topic: &topic,
+		},
+	}
+	deliveryChan := make(chan kafka.Event, 1)
+	err = p.Produce(msg, deliveryChan)
+	require.Error(t, err)
+	require.EqualError(t, err, "Local: Invalid argument or configuration")
+
+	select {
+	case <-deliveryChan:
+		assert.Fail(t, "there should be no events in the deliveryChan")
+	case <-time.After(1 * time.Second):
+		// assume there is no event
+	}
+
+	spans := mt.FinishedSpans()
+	assert.Len(t, spans, 1)
+}

contrib/confluentinc/confluent-kafka-go/kafka/kafka.go

@@ -334,31 +334,43 @@ func (p *Producer) Close() {
 func (p *Producer) Produce(msg *kafka.Message, deliveryChan chan kafka.Event) error {
 	span := p.startSpan(msg)
 
+	var errChan chan error
+
 	// if the user has selected a delivery channel, we will wrap it and
-	// wait for the delivery event to finish the span
+	// wait for the delivery event to finish the span.
+	// in case the Produce call returns an error, we won't receive anything
+	// in the deliveryChan, so we use errChan to make sure this goroutine exits.
 	if deliveryChan != nil {
+		errChan = make(chan error, 1)
 		oldDeliveryChan := deliveryChan
 		deliveryChan = make(chan kafka.Event)
 		go func() {
 			var err error
-			evt := <-deliveryChan
-			if msg, ok := evt.(*kafka.Message); ok {
-				// delivery errors are returned via TopicPartition.Error
-				err = msg.TopicPartition.Error
-				trackProduceOffsets(p.cfg.dataStreamsEnabled, msg, err)
+			select {
+			case evt := <-deliveryChan:
+				if msg, ok := evt.(*kafka.Message); ok {
+					// delivery errors are returned via TopicPartition.Error
+					err = msg.TopicPartition.Error
+					trackProduceOffsets(p.cfg.dataStreamsEnabled, msg, err)
+				}
+				oldDeliveryChan <- evt
+
+			case e := <-errChan:
+				err = e
 			}
 			span.Finish(tracer.WithError(err))
-			oldDeliveryChan <- evt
 		}()
 	}
 
 	setProduceCheckpoint(p.cfg.dataStreamsEnabled, p.libraryVersion, msg)
 	err := p.Producer.Produce(msg, deliveryChan)
-	// with no delivery channel or enqueue error, finish immediately
-	if err != nil || deliveryChan == nil {
-		span.Finish(tracer.WithError(err))
+	if err != nil {
+		if deliveryChan != nil {
+			errChan <- err
+		} else {
+			span.Finish(tracer.WithError(err))
+		}
 	}
-
 	return err
 }
 

contrib/confluentinc/confluent-kafka-go/kafka/kafka_test.go

@@ -23,6 +23,7 @@ import (
 	"github.com/confluentinc/confluent-kafka-go/kafka"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"go.uber.org/goleak"
 )
 
 var (
@@ -394,3 +395,40 @@ func TestNamingSchema(t *testing.T) {
 	}
 	namingschematest.NewKafkaTest(genSpans)(t)
 }
+
+// Test we don't leak goroutines and properly close the span when Produce returns an error.
+func TestProduceError(t *testing.T) {
+	defer goleak.VerifyNone(t)
+
+	mt := mocktracer.Start()
+	defer mt.Stop()
+
+	// first write a message to the topic
+	p, err := NewProducer(&kafka.ConfigMap{
+		"bootstrap.servers":   "127.0.0.1:9092",
+		"go.delivery.reports": true,
+	})
+	require.NoError(t, err)
+	defer p.Close()
+
+	topic := ""
+	msg := &kafka.Message{
+		TopicPartition: kafka.TopicPartition{
+			Topic: &topic,
+		},
+	}
+	deliveryChan := make(chan kafka.Event, 1)
+	err = p.Produce(msg, deliveryChan)
+	require.Error(t, err)
+	require.EqualError(t, err, "Local: Invalid argument or configuration")
+
+	select {
+	case <-deliveryChan:
+		assert.Fail(t, "there should be no events in the deliveryChan")
+	case <-time.After(1 * time.Second):
+		// assume there is no event
+	}
+
+	spans := mt.FinishedSpans()
+	assert.Len(t, spans, 1)
+}