Add support for concurrent dispatch and configurable prefetch

cmd/dispatcher/main.go

@@ -41,15 +41,15 @@ type envConfig struct {
 	// Should failed deliveries be requeued in the RabbitMQ?
 	Requeue bool `envconfig:"REQUEUE" default:"false" required:"true"`
 
+	// Number of concurrent messages in flight
+	PrefetchCount int           `envconfig:"PREFETCH_COUNT" default:"10" required:"false"`
 	Retry         int           `envconfig:"RETRY" required:"false"`
 	BackoffPolicy string        `envconfig:"BACKOFF_POLICY" required:"false"`
 	BackoffDelay  time.Duration `envconfig:"BACKOFF_DELAY" required:"false"`
 }
 
 const (
 	defaultBackoffDelay = 50 * time.Millisecond
-	defaultPrefetch     = 1
-	defaultPrefetchSize = 0
 )
 
 func main() {
@@ -99,15 +99,23 @@ func main() {
 	}()
 
 	err = channel.Qos(
-		defaultPrefetch,     // prefetch count
-		defaultPrefetchSize, // prefetch size
-		false,               // global
+		env.PrefetchCount, // prefetch count
+		0,                 // prefetch size
+		false,             // global
 	)
 	if err != nil {
 		logging.FromContext(ctx).Fatal("Failed to create QoS: ", err)
 	}
 
-	d := dispatcher.NewDispatcher(env.BrokerIngressURL, env.SubscriberURL, env.Requeue, env.Retry, backoffDelay, backoffPolicy)
+	d := &dispatcher.Dispatcher{
+		BrokerIngressURL: env.BrokerIngressURL,
+		SubscriberURL:    env.SubscriberURL,
+		Requeue:          env.Requeue,
+		MaxRetries:       env.Retry,
+		BackoffDelay:     backoffDelay,
+		BackoffPolicy:    backoffPolicy,
+		WorkerCount:      env.PrefetchCount,
+	}
 	if err := d.ConsumeFromQueue(ctx, channel, env.QueueName); err != nil {
 		// ignore ctx cancelled and channel closed errors
 		if errors.Is(err, context.Canceled) || errors.Is(err, amqperr.ErrClosed) {

pkg/dispatcher/dispatcher.go

@@ -19,6 +19,7 @@ package dispatcher
 import (
 	"context"
 	"encoding/json"
+	"sync"
 	"time"
 
 	"github.com/NeowayLabs/wabbit"
@@ -37,27 +38,17 @@ const (
 )
 
 type Dispatcher struct {
-	brokerIngressURL string
-	subscriberURL    string
+	BrokerIngressURL string
+	SubscriberURL    string
 
 	// Upon failure to deliver to sink, should the RabbitMQ messages be requeued.
-	// For example, if the DeadLetterSink has been configured, then do not requeue.
-	requeue bool
+	// For example, if the DeadLetterSink has been configured, then do not Requeue.
+	Requeue bool
 
-	maxRetries    int
-	backoffDelay  time.Duration
-	backoffPolicy eventingduckv1.BackoffPolicyType
-}
-
-func NewDispatcher(brokerIngressURL, subscriberURL string, requeue bool, maxRetries int, backoffDelay time.Duration, backoffPolicy eventingduckv1.BackoffPolicyType) *Dispatcher {
-	return &Dispatcher{
-		brokerIngressURL: brokerIngressURL,
-		subscriberURL:    subscriberURL,
-		requeue:          requeue,
-		maxRetries:       maxRetries,
-		backoffDelay:     backoffDelay,
-		backoffPolicy:    backoffPolicy,
-	}
+	MaxRetries    int
+	BackoffDelay  time.Duration
+	BackoffPolicy eventingduckv1.BackoffPolicyType
+	WorkerCount   int
 }
 
 // ConsumeFromQueue consumes messages from the given message channel and queue.
@@ -83,76 +74,31 @@ func (d *Dispatcher) ConsumeFromQueue(ctx context.Context, channel wabbit.Channe
 	}
 
 	logging.FromContext(ctx).Info("rabbitmq receiver started, exit with CTRL+C")
-	logging.FromContext(ctx).Infow("Starting to process messages", zap.String("queue", queueName))
+	logging.FromContext(ctx).Infow("Starting to process messages", zap.String("queue", queueName), zap.Int("workers", d.WorkerCount))
+
+	wg := &sync.WaitGroup{}
+	wg.Add(d.WorkerCount)
+	workerQueue := make(chan wabbit.Delivery, d.WorkerCount)
 
+	for i := 0; i < d.WorkerCount; i++ {
+		go d.dispatch(ctx, wg, workerQueue, ceClient)
+	}
 	for {
 		select {
 		case <-ctx.Done():
-			logging.FromContext(ctx).Info("context done, stopping message consumer")
+			logging.FromContext(ctx).Info("context done, stopping message consumers")
+			close(workerQueue)
+			wg.Wait()
 			return ctx.Err()
 
 		case msg, ok := <-msgs:
 			if !ok {
-				logging.FromContext(ctx).Warn("message channel closed, stopping message consumer")
+				logging.FromContext(ctx).Warn("message channel closed, stopping message consumers")
+				close(workerQueue)
+				wg.Wait()
 				return amqperr.ErrClosed
 			}
-
-			event := cloudevents.NewEvent()
-			err := json.Unmarshal(msg.Body(), &event)
-			if err != nil {
-				logging.FromContext(ctx).Warn("failed to unmarshal event (NACK-ing and not re-queueing): ", err)
-				err = msg.Nack(ackMultiple, false) // do not requeue
-				if err != nil {
-					logging.FromContext(ctx).Warn("failed to NACK event: ", err)
-				}
-				continue
-			}
-
-			logging.FromContext(ctx).Debugf("Got event as: %+v", event)
-			ctx = cloudevents.ContextWithTarget(ctx, d.subscriberURL)
-
-			// Our dispatcher uses Retries, but cloudevents is the max total tries. So we need
-			// to adjust to initial + retries.
-			// TODO: What happens if I specify 0 to cloudevents. Does it not even retry.
-			retryCount := d.maxRetries
-			if d.backoffPolicy == eventingduckv1.BackoffPolicyLinear {
-				ctx = cloudevents.ContextWithRetriesLinearBackoff(ctx, d.backoffDelay, retryCount)
-			} else {
-				ctx = cloudevents.ContextWithRetriesExponentialBackoff(ctx, d.backoffDelay, retryCount)
-			}
-
-			response, result := ceClient.Request(ctx, event)
-			if !isSuccess(ctx, result) {
-				logging.FromContext(ctx).Warnf("Failed to deliver to %q requeue: %v", d.subscriberURL, d.requeue)
-				err = msg.Nack(ackMultiple, d.requeue)
-				if err != nil {
-					logging.FromContext(ctx).Warn("failed to NACK event: ", err)
-				}
-				continue
-			}
-
-			logging.FromContext(ctx).Debugf("Got Response: %+v", response)
-			if response != nil {
-				logging.FromContext(ctx).Infof("Sending an event: %+v", response)
-				ctx = cloudevents.ContextWithTarget(ctx, d.brokerIngressURL)
-				backoffDelay := 50 * time.Millisecond
-				// Use the retries so we can just parse out the results in a common way.
-				cloudevents.ContextWithRetriesExponentialBackoff(ctx, backoffDelay, 1)
-				result := ceClient.Send(ctx, *response)
-				if !isSuccess(ctx, result) {
-					logging.FromContext(ctx).Warnf("Failed to deliver to %q requeue: %v", d.brokerIngressURL, d.requeue)
-					err = msg.Nack(ackMultiple, d.requeue) // not multiple
-					if err != nil {
-						logging.FromContext(ctx).Warn("failed to NACK event: ", err)
-					}
-					continue
-				}
-			}
-
-			err = msg.Ack(ackMultiple)
-			if err != nil {
-				logging.FromContext(ctx).Warn("failed to ACK event: ", err)
-			}
+			workerQueue <- msg
 		}
 	}
 }
@@ -172,10 +118,67 @@ func isSuccess(ctx context.Context, result protocol.Result) bool {
 				return false
 			}
 		}
-		logging.FromContext(ctx).Warnf("Invalid result type, not HTTP Result")
+		logging.FromContext(ctx).Warnf("Invalid result type, not HTTP Result: %v", retriesResult.Result)
 		return false
 	}
 
 	logging.FromContext(ctx).Warnf("Invalid result type, not RetriesResult")
 	return false
 }
+
+func (d *Dispatcher) dispatch(ctx context.Context, wg *sync.WaitGroup, queue <-chan wabbit.Delivery, ceClient cloudevents.Client) {
+	defer wg.Done()
+
+	for msg := range queue {
+		event := cloudevents.NewEvent()
+		err := json.Unmarshal(msg.Body(), &event)
+		if err != nil {
+			logging.FromContext(ctx).Warn("failed to unmarshal event (NACK-ing and not re-queueing): ", err)
+			err = msg.Nack(ackMultiple, false) // do not requeue
+			if err != nil {
+				logging.FromContext(ctx).Warn("failed to NACK event: ", err)
+			}
+			continue
+		}
+
+		logging.FromContext(ctx).Debugf("Got event as: %+v", event)
+		ctx = cloudevents.ContextWithTarget(ctx, d.SubscriberURL)
+
+		if d.BackoffPolicy == eventingduckv1.BackoffPolicyLinear {
+			ctx = cloudevents.ContextWithRetriesLinearBackoff(ctx, d.BackoffDelay, d.MaxRetries)
+		} else {
+			ctx = cloudevents.ContextWithRetriesExponentialBackoff(ctx, d.BackoffDelay, d.MaxRetries)
+		}
+
+		response, result := ceClient.Request(ctx, event)
+		if !isSuccess(ctx, result) {
+			logging.FromContext(ctx).Warnf("Failed to deliver to %q requeue: %v", d.SubscriberURL, d.Requeue)
+			err = msg.Nack(ackMultiple, d.Requeue)
+			if err != nil {
+				logging.FromContext(ctx).Warn("failed to NACK event: ", err)
+			}
+			continue
+		}
+
+		logging.FromContext(ctx).Debugf("Got Response: %+v", response)
+		if response != nil {
+			logging.FromContext(ctx).Infof("Sending an event: %+v", response)
+			ctx = cloudevents.ContextWithTarget(ctx, d.BrokerIngressURL)
+			cloudevents.ContextWithRetriesExponentialBackoff(ctx, d.BackoffDelay, d.MaxRetries)
+			result := ceClient.Send(ctx, *response)
+			if !isSuccess(ctx, result) {
+				logging.FromContext(ctx).Warnf("Failed to deliver to %q requeue: %v", d.BrokerIngressURL, d.Requeue)
+				err = msg.Nack(ackMultiple, d.Requeue) // not multiple
+				if err != nil {
+					logging.FromContext(ctx).Warn("failed to NACK event: ", err)
+				}
+				continue
+			}
+		}
+
+		err = msg.Ack(ackMultiple)
+		if err != nil {
+			logging.FromContext(ctx).Warn("failed to ACK event: ", err)
+		}
+	}
+}

pkg/dispatcher/dispatcher_test.go

@@ -35,6 +35,7 @@ import (
 
 	ce "github.com/cloudevents/sdk-go/v2/event"
 	"github.com/google/go-cmp/cmp"
+	"github.com/google/go-cmp/cmp/cmpopts"
 	eventingduckv1 "knative.dev/eventing/pkg/apis/duck/v1"
 )
 
@@ -128,15 +129,11 @@ func TestFailToConsume(t *testing.T) {
 	var buf bytes.Buffer
 	log.SetOutput(&buf)
 
-	backoffDelay, err := time.ParseDuration("1s")
-	if err != nil {
-		t.Error("Failed to parse duration: ", err)
-	}
 	ch, _, err := createRabbitAndQueue()
 	if err != nil {
 		t.Error("Failed to create rabbit and queue")
 	}
-	d := NewDispatcher("", "", false, 1, backoffDelay, eventingduckv1.BackoffPolicyExponential)
+	d := &Dispatcher{}
 	err = d.ConsumeFromQueue(context.TODO(), ch, "nosuchqueue")
 	if err == nil {
 		t.Fatal("Did not fail to consume.", err)
@@ -322,7 +319,15 @@ func TestEndToEnd(t *testing.T) {
 				backoffPolicy = eventingduckv1.BackoffPolicyLinear
 
 			}
-			d := NewDispatcher(broker.URL, subscriber.URL, tc.requeue, tc.maxRetries, backoffDelay, backoffPolicy)
+			d := &Dispatcher{
+				BrokerIngressURL: broker.URL,
+				SubscriberURL:    subscriber.URL,
+				Requeue:          tc.requeue,
+				MaxRetries:       tc.maxRetries,
+				BackoffDelay:     backoffDelay,
+				BackoffPolicy:    backoffPolicy,
+				WorkerCount:      1,
+			}
 
 			ctx, cancel := context.WithCancel(context.Background())
 			defer cancel()
@@ -365,10 +370,8 @@ func TestEndToEnd(t *testing.T) {
 			if subscriberHandler.getReceivedCount() != tc.subscriberReceiveCount {
 				t.Errorf("subscriber got %d events, wanted %d", subscriberHandler.getReceivedCount(), tc.subscriberReceiveCount)
 			} else {
-				for i := range tc.expectedSubscriberBodies {
-					if diff := cmp.Diff(tc.expectedSubscriberBodies[i], subscriberHandler.getBodies()[i]); diff != "" {
-						t.Error("unexpected subscriber diff (-want, +got) = ", diff)
-					}
+				if diff := cmp.Diff(tc.expectedSubscriberBodies, subscriberHandler.getBodies(), cmpopts.SortSlices(stringSort)); diff != "" {
+					t.Error("unexpected subscriber diff (-want, +got) = ", diff)
 				}
 			}
 			if brokerHandler.getReceivedCount() != tc.brokerReceiveCount {
@@ -440,3 +443,6 @@ func createEvent(data string) ce.Event {
 	event.SetData(cloudevents.ApplicationJSON, data)
 	return event
 }
+func stringSort(x, y string) bool {
+	return x < y
+}