fix(test): remove time.Sleep from tests for deterministic synchronization

pkg/api/message/export_test.go

@@ -1,39 +1,32 @@
 package message
 
 import (
-	"context"
-	"time"
-
 	"github.com/xmtp/xmtpd/pkg/db"
 )
 
-// AwaitCursor blocks until the subscribe worker has polled past all sequence IDs in vc.
-// Only compiled during testing (export_test.go pattern).
-func (s *Service) AwaitCursor(ctx context.Context, vc db.VectorClock) error {
-	const checkInterval = 5 * time.Millisecond
-	ticker := time.NewTicker(checkInterval)
-	defer ticker.Stop()
+// DispatchedMet reports whether the subscribe worker has *dispatched* every
+// sequence ID in vc to its listeners. This is stronger than "polled past" —
+// a true return guarantees the start() loop has already handed those rows
+// off, so any listener registered after this call will not retroactively
+// receive them. Tests use this predicate with require.Eventually to pre-seed
+// envelopes before opening a stream without racing the dispatch loop.
+func (s *Service) DispatchedMet(vc db.VectorClock) bool {
+	return s.subscribeWorker.dispatchedMet(vc)
+}
 
-	for {
-		if s.subscribeWorker.subscriptions.cursorMet(vc) {
-			return nil
-		}
-		select {
-		case <-ctx.Done():
-			return ctx.Err()
-		case <-ticker.C:
-		}
-	}
+// GlobalListenerCount returns the number of global listeners currently
+// registered with the subscribe worker. Tests poll this (via require.Eventually)
+// to wait for proof that a server-side listener is registered before
+// triggering the code under test — no time.Sleep needed.
+func (s *Service) GlobalListenerCount() int {
+	return s.subscribeWorker.countGlobalListeners()
 }
 
-func (s *subscriptionHandler) cursorMet(vc db.VectorClock) bool {
-	s.Lock()
-	defer s.Unlock()
-	for nodeID, minSeq := range vc {
-		poller, ok := s.subs[nodeID]
-		if !ok || uint64(poller.sub.LastSeen()) < minSeq {
-			return false
-		}
-	}
-	return true
+func (s *subscribeWorker) countGlobalListeners() int {
+	count := 0
+	s.globalListeners.Range(func(_, _ any) bool {
+		count++
+		return true
+	})
+	return count
 }

pkg/api/message/publish_test.go

@@ -596,13 +596,16 @@ func TestPublishEnvelopeBatchPublishNoPartialError(t *testing.T) {
 	require.Nil(t, resp)
 	require.Contains(t, err.Error(), "published via the blockchain")
 
-	// give this some time to process just in case
-	time.Sleep(100 * time.Millisecond)
-
-	envs, err := queries.New(suite.DB).
-		SelectGatewayEnvelopesUnfiltered(context.Background(), queries.SelectGatewayEnvelopesUnfilteredParams{})
-	require.NoError(t, err)
-	require.Empty(t, envs)
+	// Assert that no envelopes ever land in the DB. require.Never fails fast
+	// if one does, instead of always waiting 100ms.
+	require.Never(t, func() bool {
+		envs, err := queries.New(suite.DB).
+			SelectGatewayEnvelopesUnfiltered(
+				context.Background(),
+				queries.SelectGatewayEnvelopesUnfilteredParams{},
+			)
+		return err != nil || len(envs) > 0
+	}, 100*time.Millisecond, 10*time.Millisecond)
 }
 
 func TestPublishEnvelopeBalanceEnforcement(t *testing.T) {

pkg/api/message/subscribe_test.go

@@ -14,7 +14,6 @@ import (
 	"connectrpc.com/connect"
 	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/stretchr/testify/require"
-	"github.com/xmtp/xmtpd/pkg/api/message"
 	"github.com/xmtp/xmtpd/pkg/constants"
 	"github.com/xmtp/xmtpd/pkg/db"
 	"github.com/xmtp/xmtpd/pkg/db/queries"
@@ -112,9 +111,9 @@ func insertInitialRows(t *testing.T, suite *testUtilsApi.APIServerTestSuite) {
 	})
 	// Wait until the subscribe worker has polled past the inserted rows so that
 	// a subsequent subscription with LastSeen=nil won't see them.
-	ctx, cancel := context.WithTimeout(t.Context(), 500*time.Millisecond)
-	defer cancel()
-	require.NoError(t, suite.MessageService.AwaitCursor(ctx, db.VectorClock{100: 1, 200: 1}))
+	require.Eventually(t, func() bool {
+		return suite.MessageService.DispatchedMet(db.VectorClock{100: 1, 200: 1})
+	}, 500*time.Millisecond, 5*time.Millisecond)
 }
 
 func insertAdditionalRows(t *testing.T, store *sql.DB, notifyChan ...chan bool) {
@@ -616,8 +615,12 @@ func TestSubscribeCatchUpSkewedOriginators(t *testing.T) {
 	// Populate the database.
 	saveEnvelopes(t, server.DB, sourceEnvelopes)
 
-	// Let the subscribeWorker's catch up.
-	time.Sleep(4 * message.SubscribeWorkerPollTime)
+	// Block until the subscribeWorker has polled past the last inserted sequence ID.
+	require.Eventually(t, func() bool {
+		return server.MessageService.DispatchedMet(
+			db.VectorClock{heavyOriginatorID: uint64(heavyMsgCount)},
+		)
+	}, 5*time.Second, 5*time.Millisecond)
 
 	ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second)
 	defer cancel()
@@ -694,7 +697,6 @@ func TestSubscribeAll(t *testing.T) {
 		minEnvelopes = 10
 		maxEnvelopes = 20
 
-		insertDelay  = 100 * time.Millisecond
 		envelopeList = flattenEnvelopeMap(
 			generateEnvelopes(
 				t,
@@ -716,12 +718,12 @@ func TestSubscribeAll(t *testing.T) {
 	require.NoError(t, err)
 
 	var (
-		received = 0
+		received atomic.Int64
 		streamWG sync.WaitGroup
 	)
 
 	streamWG.Go(func() {
-		for received < total {
+		for received.Load() < int64(total) {
 			ok := stream.Receive()
 			if !ok {
 				break
@@ -730,23 +732,25 @@ func TestSubscribeAll(t *testing.T) {
 			n := len(stream.Msg().GetEnvelopes())
 			t.Logf("stream produced %v envelopes", n)
 
-			received += n
+			received.Add(int64(n))
 		}
 
 		cancel()
 	})
 
-	// Wait a bit - then start inserting envelopes. Make sure these are streamed.
-	time.Sleep(insertDelay)
+	// Wait until the server has registered the listener before inserting, so
+	// no inserts race the listener registration.
+	require.Eventually(t, func() bool {
+		return server.MessageService.GlobalListenerCount() >= 1
+	}, 5*time.Second, 5*time.Millisecond)
 
 	for _, env := range envelopeList {
 		testutils.InsertGatewayEnvelopes(t, server.DB, []queries.InsertGatewayEnvelopeV3Params{env})
-		time.Sleep(insertDelay)
 	}
 
 	streamWG.Wait()
 
-	require.Equal(t, total, received)
+	require.Equal(t, int64(total), received.Load())
 }
 
 func readOriginatorsStream(
@@ -1080,11 +1084,11 @@ func TestOriginatorParity_SkewedPagination(t *testing.T) {
 			)
 			saveEnvelopes(t, server.DB, sourceEnvelopes)
 
-			awaitCtx, awaitCancel := context.WithTimeout(t.Context(), 5*time.Second)
-			defer awaitCancel()
-			require.NoError(t, server.MessageService.AwaitCursor(
-				awaitCtx, db.VectorClock{heavyOriginatorID: uint64(heavyMsgCount)},
-			))
+			require.Eventually(t, func() bool {
+				return server.MessageService.DispatchedMet(
+					db.VectorClock{heavyOriginatorID: uint64(heavyMsgCount)},
+				)
+			}, 5*time.Second, 5*time.Millisecond)
 
 			ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second)
 			defer cancel()
@@ -1273,9 +1277,9 @@ func TestOriginatorParity_VariableVolume(t *testing.T) {
 			for nodeID, envs := range sourceEnvelopes {
 				expectedVC[uint32(nodeID)] = uint64(len(envs))
 			}
-			awaitCtx, awaitCancel := context.WithTimeout(t.Context(), 5*time.Second)
-			defer awaitCancel()
-			require.NoError(t, server.MessageService.AwaitCursor(awaitCtx, expectedVC))
+			require.Eventually(t, func() bool {
+				return server.MessageService.DispatchedMet(expectedVC)
+			}, 5*time.Second, 5*time.Millisecond)
 
 			ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second)
 			defer cancel()
@@ -1335,52 +1339,63 @@ func TestSubscribeAll_StreamsOnlyNewMessages(t *testing.T) {
 			topic.TopicKindGroupMessagesV1,
 			fmt.Appendf(nil, "generic-topic-%v", rand.Int()),
 		)
-
-		insertDelay = 100 * time.Millisecond
 	)
 
-	// Envelope data.
+	// generateEnvelopes returns perNode envelopes per originator, so flattening
+	// gives us 2*perNode total. Slice by exact bounds so the post-subscribe
+	// insert count is exactly streamSize — the final equality check depends on it.
 	var (
 		initialBatchSize = 5
 		streamSize       = 5
-		totalMessages    = initialBatchSize + streamSize
+		perNode          = initialBatchSize + streamSize
 
 		sourceEnvelopes = flattenEnvelopeMap(
 			generateEnvelopes(
 				t,
 				nodeIDs,
-				totalMessages,
-				totalMessages, // Let's get exactly N messages.
+				perNode,
+				perNode, // Exactly perNode per node.
 				payerID,
 				subTopic,
 			))
 
 		initialBatch = sourceEnvelopes[:initialBatchSize]
-		streamBatch  = sourceEnvelopes[initialBatchSize:]
+		streamBatch  = sourceEnvelopes[initialBatchSize : initialBatchSize+streamSize]
 	)
 	defer cancel()
 
-	// Pre-seed envelopes in the DB.
-	// These should NOT get picked up by the stream.
+	// Pre-seed envelopes in the DB. These should NOT get picked up by the
+	// stream because the subscribe worker marks them as known before the
+	// stream subscription is registered.
 	for _, env := range initialBatch {
 		testutils.InsertGatewayEnvelopes(t, server.DB, []queries.InsertGatewayEnvelopeV3Params{env})
 	}
 
-	// Add a delay so the subscribe worker picks pre-seeded envelopes as known before the streaming started.
-	time.Sleep(insertDelay)
+	// Block until the subscribe worker has polled past every pre-seeded row.
+	preSeedVC := make(db.VectorClock)
+	for _, env := range initialBatch {
+		nodeID := uint32(env.OriginatorNodeID)
+		seq := uint64(env.OriginatorSequenceID)
+		if cur, ok := preSeedVC[nodeID]; !ok || seq > cur {
+			preSeedVC[nodeID] = seq
+		}
+	}
+	require.Eventually(t, func() bool {
+		return server.MessageService.DispatchedMet(preSeedVC)
+	}, 5*time.Second, 5*time.Millisecond)
 
 	// Start a subscriber stream.
 	req := &message_api.SubscribeAllEnvelopesRequest{}
 	stream, err := server.ClientNotification.SubscribeAllEnvelopes(ctx, connect.NewRequest(req))
 	require.NoError(t, err)
 
 	var (
-		received = 0
+		received atomic.Int64
 		streamWG sync.WaitGroup
 	)
 
 	streamWG.Go(func() {
-		for received < streamSize {
+		for received.Load() < int64(streamSize) {
 			ok := stream.Receive()
 			if !ok {
 				break
@@ -1389,21 +1404,23 @@ func TestSubscribeAll_StreamsOnlyNewMessages(t *testing.T) {
 			n := len(stream.Msg().GetEnvelopes())
 			t.Logf("stream produced %v envelopes", n)
 
-			received += n
+			received.Add(int64(n))
 		}
 
 		cancel()
 	})
 
-	// Wait a bit - then start inserting envelopes. These should in fact be streamed.
-	time.Sleep(insertDelay)
+	// Wait until the server has registered the listener before inserting, so
+	// no inserts race the listener registration.
+	require.Eventually(t, func() bool {
+		return server.MessageService.GlobalListenerCount() >= 1
+	}, 5*time.Second, 5*time.Millisecond)
 
 	for _, env := range streamBatch {
 		testutils.InsertGatewayEnvelopes(t, server.DB, []queries.InsertGatewayEnvelopeV3Params{env})
-		time.Sleep(insertDelay)
 	}
 
 	streamWG.Wait()
 
-	require.Equal(t, streamSize, received)
+	require.Equal(t, int64(streamSize), received.Load())
 }