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

pkg/api/message/export_test.go

@@ -7,15 +7,19 @@ import (
 	"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).
+// AwaitCursor blocks until the subscribe worker has *dispatched* every
+// sequence ID in vc to its listeners. This is stronger than "polled past" —
+// it guarantees the start() loop has already handed those rows off, so any
+// listener registered after this call will not retroactively receive them.
+// Tests rely on this guarantee to pre-seed envelopes before opening a stream
+// without racing the worker's dispatch loop.
 func (s *Service) AwaitCursor(ctx context.Context, vc db.VectorClock) error {
 	const checkInterval = 5 * time.Millisecond
 	ticker := time.NewTicker(checkInterval)
 	defer ticker.Stop()
 
 	for {
-		if s.subscribeWorker.subscriptions.cursorMet(vc) {
+		if s.subscribeWorker.dispatchedMet(vc) {
 			return nil
 		}
 		select {
@@ -26,14 +30,33 @@ func (s *Service) AwaitCursor(ctx context.Context, vc db.VectorClock) error {
 	}
 }
 
-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
+// AwaitGlobalListeners blocks until at least n global listeners are registered
+// with the subscribe worker. Tests use this to eliminate the races between
+// opening a SubscribeAllEnvelopes stream and inserting envelopes — instead of
+// sleeping, they wait for proof that the server-side listener is registered
+// before triggering the code under test.
+func (s *Service) AwaitGlobalListeners(ctx context.Context, n int) error {
+	const checkInterval = 5 * time.Millisecond
+	ticker := time.NewTicker(checkInterval)
+	defer ticker.Stop()
+
+	for {
+		if s.subscribeWorker.countGlobalListeners() >= n {
+			return nil
+		}
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case <-ticker.C:
 		}
 	}
-	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"
@@ -616,8 +615,18 @@ 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.
+	{
+		awaitCtx, awaitCancel := context.WithTimeout(t.Context(), 5*time.Second)
+		require.NoError(
+			t,
+			server.MessageService.AwaitCursor(
+				awaitCtx,
+				db.VectorClock{heavyOriginatorID: uint64(heavyMsgCount)},
+			),
+		)
+		awaitCancel()
+	}
 
 	ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second)
 	defer cancel()
@@ -694,7 +703,6 @@ func TestSubscribeAll(t *testing.T) {
 		minEnvelopes = 10
 		maxEnvelopes = 20
 
-		insertDelay  = 100 * time.Millisecond
 		envelopeList = flattenEnvelopeMap(
 			generateEnvelopes(
 				t,
@@ -716,12 +724,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 +738,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.
+	awaitCtx, awaitCancel := context.WithTimeout(t.Context(), 5*time.Second)
+	require.NoError(t, server.MessageService.AwaitGlobalListeners(awaitCtx, 1))
+	awaitCancel()
 
 	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(
@@ -1335,52 +1345,69 @@ 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 with
+	// 2 nodes we get 2 * perNode total envelopes. We only insert the first
+	// `initialBatchSize` as pre-seeds and the next `streamSize` as the new
+	// batch the stream should observe; the remainder are generated but never
+	// inserted. Slicing by exact bounds (rather than [initialBatchSize:])
+	// keeps the post-subscribe insert count to exactly streamSize — otherwise
+	// the stream races past streamSize and the final equality check fails.
 	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
+		}
+	}
+	awaitCtx, awaitCancel := context.WithTimeout(t.Context(), 5*time.Second)
+	require.NoError(t, server.MessageService.AwaitCursor(awaitCtx, preSeedVC))
+	awaitCancel()
 
 	// 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 +1416,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.
+	listenerCtx, listenerCancel := context.WithTimeout(t.Context(), 5*time.Second)
+	require.NoError(t, server.MessageService.AwaitGlobalListeners(listenerCtx, 1))
+	listenerCancel()
 
 	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())
 }