commit offsets and wait for consumers to stop (#98)

Author: travis-minke-sap

components/dispatcher/cmd/main.go

@@ -172,9 +172,11 @@ func main() {
 	healthServer.SetAlive(true)
 	healthServer.SetDispatcherReady(true)
 
+	// Start The Controllers
 	logger.Info("Starting controllers.")
 	kncontroller.StartAll(stopCh, controllers[:]...)
 
+	// Block On Signal Handler Stop Channel
 	<-stopCh
 
 	// Reset The Liveness and Readiness Flags In Preparation For Shutdown

components/dispatcher/internal/controller/kafkachannel_test.go

@@ -137,6 +137,9 @@ func TestAllCases(t *testing.T) {
 			KafkaClientSet:       kafkaClient,
 		}
 	}))
+
+	// Pause to let async go processes finish logging :(
+	time.Sleep(1 * time.Second)
 }
 
 func NewTestDispatcher(t *testing.T, channelKey string) *dispatcher.Dispatcher {

components/dispatcher/internal/dispatcher/dispatcher.go

@@ -40,6 +40,7 @@ type ConsumerOffset struct {
 	lastOffsetCommit time.Time
 	offsets          map[int32]kafka.Offset
 	stopCh           chan bool
+	stoppedCh        chan bool
 }
 
 // Define a Dispatcher Struct to hold Dispatcher Config and dispatcher implementation details
@@ -62,17 +63,19 @@ func NewDispatcher(dispatcherConfig DispatcherConfig) *Dispatcher {
 	return dispatcher
 }
 
-// Close The Running Consumers Connections
+// Stop All Consumers
 func (d *Dispatcher) StopConsumers() {
 	for subscription, _ := range d.consumers {
 		d.stopConsumer(subscription)
 	}
 }
 
+// Stop An Individual Consumer
 func (d *Dispatcher) stopConsumer(subscription Subscription) {
 	d.Logger.Info("Stopping Consumer", zap.String("GroupId", subscription.GroupId), zap.String("topic", d.Topic), zap.String("URI", subscription.URI))
-
-	d.consumers[subscription].stopCh <- true
+	consumerOffset := d.consumers[subscription]
+	consumerOffset.stopCh <- true // Send Stop Signal
+	<-consumerOffset.stoppedCh    // Wait Until Stop Completes
 	delete(d.consumers, subscription)
 }
 
@@ -94,8 +97,14 @@ func (d *Dispatcher) initConsumer(subscription Subscription) (*ConsumerOffset, e
 		return nil, err
 	}
 
-	stopCh := make(chan bool)
-	consumerOffset := ConsumerOffset{consumer: consumer, lastOffsetCommit: time.Now(), offsets: make(map[int32]kafka.Offset), stopCh: stopCh}
+	// Create The ConsumerOffset For Tracking State
+	consumerOffset := ConsumerOffset{
+		consumer:         consumer,
+		lastOffsetCommit: time.Now(),
+		offsets:          make(map[int32]kafka.Offset),
+		stopCh:           make(chan bool),
+		stoppedCh:        make(chan bool),
+	}
 
 	// Start Consuming Messages From Topic (Async)
 	go d.handleKafkaMessages(consumerOffset, subscription)
@@ -110,23 +119,26 @@ func (d *Dispatcher) handleKafkaMessages(consumerOffset ConsumerOffset, subscrip
 	// Configure The Logger
 	logger := d.Logger.With(zap.String("GroupID", subscription.GroupId))
 
-	stopped := false
 	// Message Processing Loop
+	stopped := false
 	for !stopped {
 
 		// Poll For A New Event Message Until We Get One (Timeout is how long to wait before requesting again)
 		event := consumerOffset.consumer.Poll(d.PollTimeoutMillis)
 
 		select {
-		// Handle Shutdown Case
+
+		// Non-Blocking Channel Check
 		case <-consumerOffset.stopCh:
-			stopped = true
+			stopped = true // Handle Shutdown Case - Break Out Of Loop
+
 		default:
 
+			// No Events To Process At The Moment
 			if event == nil {
+
 				// Commit Offsets If The Amount Of Time Since Last Commit Is "OffsetCommitDuration" Or More
 				currentTimeDuration := time.Now().Sub(consumerOffset.lastOffsetCommit)
-
 				if currentTimeDuration > d.OffsetCommitDurationMinimum && currentTimeDuration >= d.OffsetCommitDuration {
 					d.commitOffsets(logger, &consumerOffset)
 				}
@@ -177,12 +189,17 @@ func (d *Dispatcher) handleKafkaMessages(consumerOffset ConsumerOffset, subscrip
 		}
 	}
 
+	// Commit Offsets One Last Time
+	logger.Debug("Final Offset Commit")
+	d.commitOffsets(logger, &consumerOffset)
+
 	// Safe To Shutdown The Consumer
 	logger.Debug("Shutting Down Consumer")
 	err := consumerOffset.consumer.Close()
 	if err != nil {
 		logger.Error("Unable To Stop Consumer", zap.Error(err))
 	}
+	consumerOffset.stoppedCh <- true
 }
 
 // Store Updated Offsets For The Partition If Consumer Still Has It Assigned
@@ -213,6 +230,7 @@ func (d *Dispatcher) commitOffsets(logger *zap.Logger, consumerOffset *ConsumerO
 	consumerOffset.lastOffsetCommit = time.Now()
 }
 
+// Update The Dispatcher's Subscriptions
 func (d *Dispatcher) UpdateSubscriptions(subscriptions []Subscription) map[Subscription]error {
 
 	failedSubscriptions := make(map[Subscription]error)

components/dispatcher/internal/dispatcher/dispatcher_test.go

@@ -24,7 +24,7 @@ const (
 	testTopic                   = "TestTopic"
 	testPartition               = 33
 	testOffset                  = "latest"
-	testGroupId                 = "TestGroupId"
+	testGroupId1                = "TestGroupId1"
 	testGroupId2                = "TestGroupId2"
 	testGroupId3                = "TestGroupId3"
 	testKey                     = "TestKey"
@@ -113,7 +113,7 @@ func TestDispatcher(t *testing.T) {
 	subscriptionResults := testDispatcher.UpdateSubscriptions([]Subscription{
 		{
 			URI:     testSubscriberUri1,
-			GroupId: testGroupId,
+			GroupId: testGroupId1,
 		},
 	})
 
@@ -131,7 +131,7 @@ func TestDispatcher(t *testing.T) {
 	subscriptionResults = testDispatcher.UpdateSubscriptions([]Subscription{
 		{
 			URI:     testSubscriberUri1,
-			GroupId: testGroupId,
+			GroupId: testGroupId1,
 		},
 		{
 			URI:     testSubscriberUri2,
@@ -153,7 +153,7 @@ func TestDispatcher(t *testing.T) {
 	sendMessagesToConsumers(t, testDispatcher.consumers, testMessagesToSend)
 
 	// Wait For Consumers To Process Messages
-	waitForConsumersToProcessEvents(t, testDispatcher.consumers)
+	waitForConsumersToProcessEvents(t, testDispatcher.consumers, testMessagesToSend)
 
 	// Verify The Consumer Offset Commit Messages
 	verifyConsumerCommits(t, testDispatcher.consumers, 3)
@@ -162,7 +162,7 @@ func TestDispatcher(t *testing.T) {
 	subscriptionResults = testDispatcher.UpdateSubscriptions([]Subscription{
 		{
 			URI:     testSubscriberUri1,
-			GroupId: testGroupId,
+			GroupId: testGroupId1,
 		},
 	})
 	verifyConsumersCount(t, testDispatcher.consumers, 1)
@@ -393,21 +393,34 @@ func createMessageHeaders(context cloudevents.EventContext) []kafka.Header {
 }
 
 // Wait For The Consumers To Process All The Events
-func waitForConsumersToProcessEvents(t *testing.T, consumers map[Subscription]*ConsumerOffset) {
+func waitForConsumersToProcessEvents(t *testing.T, consumers map[Subscription]*ConsumerOffset, expectedOffset kafka.Offset) {
+
+	// Track Wait Start Time
 	startTime := time.Now()
+
+	// Loop Over All The Specified Consumers
 	for _, consumer := range consumers {
+
+		// Verified Cast Of Consumer As MockConsumer
 		mockConsumer, ok := consumer.consumer.(*MockConsumer)
 		assert.True(t, ok)
+
+		// Infinite Loop - Broken By Internal Timeout Failure
 		for {
+
+			// Get The Partition's Committed Offset (Minimal Synchronization To Avoid Deadlocking The Commit Updates!)
 			mockConsumer.offsetsMutex.Lock()
-			if mockConsumer.getCommits()[testPartition] >= testMessagesToSend+1 {
+			commitOffset := mockConsumer.getCommits()[testPartition]
+			mockConsumer.offsetsMutex.Unlock()
+
+			// Process Committed Offset - Either Stop Looking, Wait, Or Timeout
+			if commitOffset >= expectedOffset+1 {
 				break
-			} else if time.Now().Sub(startTime) > (4 * time.Second) {
+			} else if time.Now().Sub(startTime) > (5 * time.Second) {
 				assert.FailNow(t, "Timed-out Waiting For Consumers To Process Events")
 			} else {
 				time.Sleep(100 * time.Millisecond)
 			}
-			mockConsumer.offsetsMutex.Unlock()
 		}
 	}
 }