Merge pull request #133 from uw-labs/instrumented-sink-detect-context-cancel

Instrumented async message sink: do not increment error counter if the context was cancelled

Author: sbuliarca

instrumented/sink.go

@@ -2,6 +2,7 @@ package instrumented
 
 import (
 	"context"
+	"errors"
 
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/uw-labs/substrate"
@@ -57,22 +58,34 @@ func (ams *AsyncMessageSink) PublishMessages(ctx context.Context, acks chan<- su
 			case <-ctx.Done():
 				return <-errs
 			case err := <-errs:
-				if err != nil {
+				if isUnexpectedError(err) {
 					ams.counter.WithLabelValues("error", ams.topic).Inc()
 				}
 				return err
 			}
 		case <-ctx.Done():
 			return <-errs
 		case err := <-errs:
-			if err != nil {
+			if isUnexpectedError(err) {
 				ams.counter.WithLabelValues("error", ams.topic).Inc()
 			}
 			return err
 		}
 	}
 }
 
+func isUnexpectedError(err error) bool {
+	switch {
+	case err == nil:
+		return false
+	case errors.Is(err, context.Canceled):
+		// we're expecting producers to mark the stopping of producing by cancelling the context
+		return false
+	default:
+		return true
+	}
+}
+
 // Close closes the message sink
 func (ams *AsyncMessageSink) Close() error {
 	return ams.impl.Close()

instrumented/sink_test.go

@@ -9,6 +9,7 @@ import (
 	"github.com/prometheus/client_golang/prometheus"
 	dto "github.com/prometheus/client_model/go"
 	"github.com/stretchr/testify/assert"
+
 	"github.com/uw-labs/substrate"
 )
 
@@ -79,6 +80,60 @@ func TestPublishMessagesSuccessfully(t *testing.T) {
 	}
 }
 
+func TestPublishMessagesSuccessfully_WithContextError(t *testing.T) {
+	sink := AsyncMessageSink{
+		impl: &asyncMessageSinkMock{
+			publishMessageMock: func(ctx context.Context, acks chan<- substrate.Message, messages <-chan substrate.Message) error {
+				for {
+					select {
+					case <-ctx.Done():
+						return ctx.Err()
+					case msg := <-messages:
+						acks <- msg
+					}
+				}
+			},
+		},
+		counter: prometheus.NewCounterVec(
+			prometheus.CounterOpts{
+				Help: "sink_counter",
+				Name: "sink_counter",
+			}, []string{"status", "topic"}),
+		topic: "testTopic",
+	}
+
+	acks := make(chan substrate.Message)
+	messages := make(chan substrate.Message)
+
+	sinkContext, sinkCancel := context.WithCancel(context.Background())
+	defer sinkCancel()
+
+	errs := make(chan error)
+	go func() {
+		defer close(errs)
+		errs <- sink.PublishMessages(sinkContext, acks, messages)
+	}()
+
+	messages <- Message{}
+
+	for {
+		select {
+		case err := <-errs:
+			assert.True(t, errors.Is(err, context.Canceled))
+			// the error counter should not be increased when producing stops by context canceling
+			var metric dto.Metric
+			assert.NoError(t, sink.counter.WithLabelValues("error", "testTopic").Write(&metric))
+			assert.Equal(t, 0, int(*metric.Counter.Value))
+			return
+		case <-acks:
+			var metric dto.Metric
+			assert.NoError(t, sink.counter.WithLabelValues("success", "testTopic").Write(&metric))
+			assert.Equal(t, 1, int(*metric.Counter.Value))
+			sinkCancel()
+		}
+	}
+}
+
 func TestPublishMessagesWithError(t *testing.T) {
 	producingError := errors.New("message producing error")
 	sink := AsyncMessageSink{