Merge pull request #115 from uw-labs/fix_more_channel_and_goroutine_issues_in_sink

Fix more channel and goroutine issues in sink

Author: mjgarton

kafka/producer.go

@@ -9,6 +9,7 @@ import (
 	"github.com/uw-labs/substrate/internal/debug"
 	"github.com/uw-labs/substrate/internal/helper"
 	"github.com/uw-labs/substrate/internal/unwrap"
+	"golang.org/x/sync/errgroup"
 )
 
 var (
@@ -79,41 +80,60 @@ func (ams *asyncMessageSink) doPublishMessages(ctx context.Context, producer sar
 	errs := producer.Errors()
 	successes := producer.Successes()
 
-	go func() {
-		for suc := range successes {
-			msg := suc.Metadata.(substrate.Message)
-			acks <- msg
-			ams.debugger.Logf("substrate : sent ack to caller for message : %s\n", msg)
-		}
-	}()
-	for {
-		select {
-		case m := <-messages:
-			message := &sarama.ProducerMessage{
-				Topic: ams.Topic,
-			}
+	eg, ctx := errgroup.WithContext(ctx)
 
-			message.Value = sarama.ByteEncoder(m.Data())
-
-			if ams.KeyFunc != nil {
-				// Provide original user message to the partition key function.
-				unwrappedMsg := unwrap.Unwrap(m)
-				message.Key = sarama.ByteEncoder(ams.KeyFunc(unwrappedMsg))
+	eg.Go(func() error {
+		for {
+			select {
+			case suc := <-successes:
+				msg := suc.Metadata.(substrate.Message)
+				select {
+				case acks <- msg:
+					ams.debugger.Logf("substrate : sent ack to caller for message : %s\n", msg)
+				case <-ctx.Done():
+					return ctx.Err()
+				}
+			case <-ctx.Done():
+				return ctx.Err()
 			}
+		}
+	})
 
-			message.Metadata = m
+	eg.Go(func() error {
+		for {
 			select {
-			case input <- message:
+			case m := <-messages:
+				message := &sarama.ProducerMessage{
+					Topic: ams.Topic,
+				}
+
+				message.Value = sarama.ByteEncoder(m.Data())
+
+				if ams.KeyFunc != nil {
+					// Provide original user message to the partition key function.
+					unwrappedMsg := unwrap.Unwrap(m)
+					message.Key = sarama.ByteEncoder(ams.KeyFunc(unwrappedMsg))
+				}
+
+				message.Metadata = m
+				select {
+				case input <- message:
+				case <-ctx.Done():
+					return ctx.Err()
+				}
+				ams.debugger.Logf("substrate : sent to kafka : %s\n", m)
 			case <-ctx.Done():
-				return nil
+				return ctx.Err()
+			case err := <-errs:
+				return err
 			}
-			ams.debugger.Logf("substrate : sent to kafka : %s\n", m)
-		case <-ctx.Done():
-			return nil
-		case err := <-errs:
-			return err
 		}
+	})
+
+	if err := eg.Wait(); err != nil && err != context.Canceled {
+		return err
 	}
+	return nil
 }
 
 func (ams *asyncMessageSink) Status() (*substrate.Status, error) {