Make eventSequencer thread safe

Author: prestonvasquez

internal/integration/unified/client_entity.go

@@ -48,9 +48,14 @@ var securitySensitiveCommands = []string{
 // recorded. After pool initialization completes, we set eventCutoffSeq to the
 // current sequence number. Event accessors for CMAP and SDAM types then
 // filter out any events with sequence <= eventCutoffSeq.
+//
+// Sequencing is thread-safe to support concurrent operations that may generate
+// events (e.g., connection checkouts generating CMAP events).
 type eventSequencer struct {
 	counter atomic.Int64
-	cutoff  int64
+	cutoff  atomic.Int64
+
+	mu sync.RWMutex
 
 	// pool events are heterogeneous, so we track their sequence separately
 	poolSeq        []int64
@@ -59,27 +64,42 @@ type eventSequencer struct {
 
 // setCutoff marks the current sequence as the filtering cutoff point.
 func (es *eventSequencer) setCutoff() {
-	es.cutoff = es.counter.Load()
+	es.cutoff.Store(es.counter.Load())
 }
 
 // recordEvent stores the sequence number for a given event type.
 func (es *eventSequencer) recordEvent(eventType monitoringEventType) {
 	next := es.counter.Add(1)
+
+	es.mu.Lock()
 	es.seqByEventType[eventType] = append(es.seqByEventType[eventType], next)
+	es.mu.Unlock()
 }
 
 func (es *eventSequencer) recordPooledEvent() {
 	next := es.counter.Add(1)
+
+	es.mu.Lock()
 	es.poolSeq = append(es.poolSeq, next)
+	es.mu.Unlock()
 }
 
 // shouldFilter returns true if the event at the given index should be filtered.
 func (es *eventSequencer) shouldFilter(eventType monitoringEventType, index int) bool {
-	if es.cutoff == 0 {
+	cutoff := es.cutoff.Load()
+	if cutoff == 0 {
 		return false
 	}
 
-	return es.seqByEventType[eventType][index] <= es.cutoff
+	es.mu.RLock()
+	defer es.mu.RUnlock()
+
+	seqs, ok := es.seqByEventType[eventType]
+	if !ok || index < 0 || index >= len(seqs) {
+		return false
+	}
+
+	return seqs[index] <= cutoff
 }
 
 // clientEntity is a wrapper for a mongo.Client object that also holds additional information required during test
@@ -352,17 +372,43 @@ func (c *clientEntity) failedEvents() []*event.CommandFailedEvent {
 	return events
 }
 
-// filterEventsBySeq filters events by sequence number using the provided
-// sequence slice. See comments on eventSequencer for more details.
-func filterEventsBySeq[T any](c *clientEntity, events []T, seqSlice []int64) []T {
-	if c.eventSequencer.cutoff == 0 {
+// filterEventsBySeq filters events by sequence number for the given eventType.
+// See comments on eventSequencer for more details.
+func filterEventsBySeq[T any](c *clientEntity, events []T, eventType monitoringEventType) []T {
+	cutoff := c.eventSequencer.cutoff.Load()
+	if cutoff == 0 {
 		return events
 	}
 
-	var filtered []T
-	for i, evt := range events {
-		if seqSlice[i] > c.eventSequencer.cutoff {
-			filtered = append(filtered, evt)
+	// Lock order: eventProcessMu -> eventSequencer.mu (matches writers)
+	c.eventProcessMu.RLock()
+	c.eventSequencer.mu.RLock()
+
+	// Snapshot to minimize time under locks and avoid races
+	localEvents := append([]T(nil), events...)
+
+	var seqSlice []int64
+	if eventType == poolAnyEvent {
+		seqSlice = c.eventSequencer.poolSeq
+	} else {
+		seqSlice = c.eventSequencer.seqByEventType[eventType]
+	}
+
+	localSeqs := append([]int64(nil), seqSlice...)
+
+	c.eventSequencer.mu.RUnlock()
+	c.eventProcessMu.RUnlock()
+
+	// guard against index out of range.
+	n := len(localEvents)
+	if len(localSeqs) < n {
+		n = len(localSeqs)
+	}
+
+	filtered := make([]T, 0, n)
+	for i := 0; i < n; i++ {
+		if localSeqs[i] > cutoff {
+			filtered = append(filtered, localEvents[i])
 		}
 	}
 
@@ -555,8 +601,10 @@ func (c *clientEntity) processPoolEvent(evt *event.PoolEvent) {
 
 	eventType := monitoringEventTypeFromPoolEvent(evt)
 	if _, ok := c.observedEvents[eventType]; ok {
+		c.eventProcessMu.Lock()
 		c.pooled = append(c.pooled, evt)
 		c.eventSequencer.recordPooledEvent()
+		c.eventProcessMu.Unlock()
 	}
 
 	c.addEventsCount(eventType)
@@ -787,9 +835,7 @@ func awaitMinimumPoolSize(ctx context.Context, entity *clientEntity, minPoolSize
 		case <-awaitCtx.Done():
 			return fmt.Errorf("timed out waiting for client to reach minPoolSize")
 		case <-ticker.C:
-			if uint64(entity.eventsCount[connectionReadyEvent]) >= minPoolSize {
-				// Clear all CMAP and SDAM events that occurred during pool
-				// initialization.
+			if uint64(entity.getEventCount(connectionReadyEvent)) >= minPoolSize {
 				entity.eventSequencer.setCutoff()
 
 				return nil

internal/integration/unified/client_entity_test.go

@@ -30,11 +30,6 @@ func recordTopologyOpening(c *clientEntity) {
 	c.eventSequencer.recordEvent(topologyOpeningEvent)
 }
 
-func recordHeartbeatSucceeded(c *clientEntity) {
-	c.serverHeartbeatSucceeded = append(c.serverHeartbeatSucceeded, &event.ServerHeartbeatSucceededEvent{})
-	c.eventSequencer.recordEvent(serverHeartbeatSucceededEvent)
-}
-
 func Test_eventSequencer(t *testing.T) {
 	tests := []struct {
 		name           string
@@ -91,19 +86,6 @@ func Test_eventSequencer(t *testing.T) {
 				topologyOpeningEvent:          1,
 			},
 		},
-		{
-			name:        "cutoff at beginning filters nothing",
-			cutoffAfter: 0,
-			setupEvents: func(c *clientEntity) {
-				// Cutoff will be set immediately (before any events)
-				recordPoolEvent(c)
-				recordHeartbeatSucceeded(c)
-			},
-			expectedPooled: 1,
-			expectedSDAM: map[monitoringEventType]int{
-				serverHeartbeatSucceededEvent: 1,
-			},
-		},
 		{
 			name:        "cutoff after all events filters everything",
 			cutoffAfter: 3,
@@ -135,25 +117,24 @@ func Test_eventSequencer(t *testing.T) {
 			// Set cutoff if specified
 			if tt.cutoffAfter > 0 {
 				// Manually set cutoff to the specified event sequence
-				client.eventSequencer.cutoff = int64(tt.cutoffAfter)
+				client.eventSequencer.cutoff.Store(int64(tt.cutoffAfter))
 			}
 
 			// Test pool event filtering
-			filteredPool := filterEventsBySeq(client, client.pooled, client.eventSequencer.poolSeq)
+			filteredPool := filterEventsBySeq(client, client.pooled, poolAnyEvent)
 			assert.Equal(t, tt.expectedPooled, len(filteredPool), "pool events count mismatch")
 
 			// Test SDAM event filtering
 			for eventType, expectedCount := range tt.expectedSDAM {
 				var actualCount int
-				seqs := client.eventSequencer.seqByEventType[eventType]
 
 				switch eventType {
 				case serverDescriptionChangedEvent:
-					actualCount = len(filterEventsBySeq(client, client.serverDescriptionChanged, seqs))
+					actualCount = len(filterEventsBySeq(client, client.serverDescriptionChanged, serverDescriptionChangedEvent))
 				case serverHeartbeatSucceededEvent:
-					actualCount = len(filterEventsBySeq(client, client.serverHeartbeatSucceeded, seqs))
+					actualCount = len(filterEventsBySeq(client, client.serverHeartbeatSucceeded, serverHeartbeatSucceededEvent))
 				case topologyOpeningEvent:
-					actualCount = len(filterEventsBySeq(client, client.topologyOpening, seqs))
+					actualCount = len(filterEventsBySeq(client, client.topologyOpening, topologyOpeningEvent))
 				}
 
 				assert.Equal(t, expectedCount, actualCount, "%s count mismatch", eventType)
@@ -180,14 +161,14 @@ func Test_eventSequencer_setCutoff(t *testing.T) {
 	client.eventSequencer.setCutoff()
 
 	// Verify cutoff matches counter
-	assert.Equal(t, int64(2), client.eventSequencer.cutoff, "cutoff should be 2")
+	assert.Equal(t, int64(2), client.eventSequencer.cutoff.Load(), "cutoff should be 2")
 
 	// Record more events
 	recordPoolEvent(client)
 
 	// Verify counter incremented but cutoff didn't
 	assert.Equal(t, int64(3), client.eventSequencer.counter.Load(), "counter should be 3")
-	assert.Equal(t, int64(2), client.eventSequencer.cutoff, "cutoff should still be 2")
+	assert.Equal(t, int64(2), client.eventSequencer.cutoff.Load(), "cutoff should still be 2")
 }
 
 func Test_eventSequencer_shouldFilter(t *testing.T) {
@@ -245,7 +226,7 @@ func Test_eventSequencer_shouldFilter(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			es.cutoff = tt.cutoff
+			es.cutoff.Store(tt.cutoff)
 			result := es.shouldFilter(tt.eventType, tt.index)
 			assert.Equal(t, tt.expected, result, "shouldFilter result mismatch")
 		})

internal/integration/unified/event.go

@@ -37,6 +37,9 @@ const (
 	topologyDescriptionChangedEvent monitoringEventType = "TopologyDescriptionChangedEvent"
 	topologyOpeningEvent            monitoringEventType = "TopologyOpeningEvent"
 	topologyClosedEvent             monitoringEventType = "TopologyClosedEvent"
+
+	// sentinel: indicates "use pooled (CMAP) sequence".
+	poolAnyEvent monitoringEventType = "_PoolAny"
 )
 
 func monitoringEventTypeFromString(eventStr string) (monitoringEventType, bool) {

internal/integration/unified/event_verification.go

@@ -312,7 +312,7 @@ func verifyCommandEvents(ctx context.Context, client *clientEntity, expectedEven
 }
 
 func verifyCMAPEvents(client *clientEntity, expectedEvents *expectedEvents) error {
-	pooled := filterEventsBySeq(client, client.pooled, client.eventSequencer.poolSeq)
+	pooled := filterEventsBySeq(client, client.pooled, poolAnyEvent)
 	if len(expectedEvents.CMAPEvents) == 0 && len(pooled) != 0 {
 		return fmt.Errorf("expected no cmap events to be sent but got %s", stringifyEventsForClient(client))
 	}
@@ -443,7 +443,7 @@ func stringifyEventsForClient(client *clientEntity) string {
 	}
 
 	str.WriteString("\nPool Events\n\n")
-	for _, evt := range filterEventsBySeq(client, client.pooled, client.eventSequencer.poolSeq) {
+	for _, evt := range filterEventsBySeq(client, client.pooled, poolAnyEvent) {
 		str.WriteString(fmt.Sprintf("[%s] Event Type: %q\n", evt.Address, evt.Type))
 	}
 
@@ -522,13 +522,13 @@ func getNextTopologyClosedEvent(
 
 func verifySDAMEvents(client *clientEntity, expectedEvents *expectedEvents) error {
 	var (
-		changed   = filterEventsBySeq(client, client.serverDescriptionChanged, client.eventSequencer.seqByEventType[serverDescriptionChangedEvent])
-		started   = filterEventsBySeq(client, client.serverHeartbeatStartedEvent, client.eventSequencer.seqByEventType[serverHeartbeatStartedEvent])
-		succeeded = filterEventsBySeq(client, client.serverHeartbeatSucceeded, client.eventSequencer.seqByEventType[serverHeartbeatSucceededEvent])
-		failed    = filterEventsBySeq(client, client.serverHeartbeatFailedEvent, client.eventSequencer.seqByEventType[serverHeartbeatFailedEvent])
-		tchanged  = filterEventsBySeq(client, client.topologyDescriptionChanged, client.eventSequencer.seqByEventType[topologyDescriptionChangedEvent])
-		topening  = filterEventsBySeq(client, client.topologyOpening, client.eventSequencer.seqByEventType[topologyOpeningEvent])
-		tclosed   = filterEventsBySeq(client, client.topologyClosed, client.eventSequencer.seqByEventType[topologyClosedEvent])
+		changed   = filterEventsBySeq(client, client.serverDescriptionChanged, serverDescriptionChangedEvent)
+		started   = filterEventsBySeq(client, client.serverHeartbeatStartedEvent, serverHeartbeatStartedEvent)
+		succeeded = filterEventsBySeq(client, client.serverHeartbeatSucceeded, serverHeartbeatSucceededEvent)
+		failed    = filterEventsBySeq(client, client.serverHeartbeatFailedEvent, serverHeartbeatFailedEvent)
+		tchanged  = filterEventsBySeq(client, client.topologyDescriptionChanged, topologyDescriptionChangedEvent)
+		topening  = filterEventsBySeq(client, client.topologyOpening, topologyOpeningEvent)
+		tclosed   = filterEventsBySeq(client, client.topologyClosed, topologyClosedEvent)
 	)
 
 	vol := func() int {