mongodb: improving idle change stream resume (#4254)

This PR:
- revert the previous 5 minute context timeout
[change](https://github.com/PeerDB-io/peerdb/pull/4166/changes) (this
was [causing an issue](#4213)
where change stream that takes a long time to resume can stuck and retry
loop).
- setting `client.Watch` with batchSize = 0, allowing change stream to
return immediately, and the "catch-up" can be blocked on
`changestream.Next` instead of `client.Watch`
- improve handle of idle stream by:
- always checkpoint in-memory when context timed out with records
accumulated
- always persist offset to catalog when context timed out without
records accumulated, so that if pipe is paused or an error occurred, we
can resume from a recent offset
- add unit test with a mock change stream to avoid having to slow down
CI; this will also make it easier to extend unit tests for other mongodb
cdc behavior

Fixes: DBI-691

Author: jgao54

flow/connectors/mongo/cdc.go

@@ -36,6 +36,23 @@ type ChangeEvent struct {
 	ClusterTime   bson.Timestamp `bson:"clusterTime"`
 }
 
+// ChangeStream is defined as an interface, allowing tests inject mock change stream.
+type ChangeStream interface {
+	Next(ctx context.Context) bool
+	ResumeToken() bson.Raw
+	Err() error
+	Close(ctx context.Context) error
+	Current() bson.Raw
+}
+
+type changeStreamWrapper struct {
+	*mongo.ChangeStream
+}
+
+func (w *changeStreamWrapper) Current() bson.Raw {
+	return w.ChangeStream.Current
+}
+
 func (c *MongoConnector) GetTableSchema(
 	ctx context.Context,
 	_ map[string]string,
@@ -87,8 +104,7 @@ func (c *MongoConnector) SetupReplication(ctx context.Context, input *protos.Set
 	if err != nil {
 		return model.SetupReplicationResult{}, fmt.Errorf("failed to create changestream pipeline: %w", err)
 	}
-
-	changeStream, err := createChangeStream(c.client, ctx, pipeline, changeStreamOpts)
+	changeStream, err := c.createChangeStream(ctx, pipeline, changeStreamOpts)
 	if err != nil {
 		return model.SetupReplicationResult{}, fmt.Errorf("failed to start change stream for storing initial resume token: %w", err)
 	}
@@ -107,7 +123,7 @@ func (c *MongoConnector) SetupReplication(ctx context.Context, input *protos.Set
 			}
 		}
 	}
-	err = c.SetLastOffset(ctx, input.FlowJobName, model.CdcCheckpoint{
+	err = c.metadataStore.SetLastOffset(ctx, input.FlowJobName, model.CdcCheckpoint{
 		Text: base64.StdEncoding.EncodeToString(resumeToken),
 	})
 	if err != nil {
@@ -137,7 +153,12 @@ func (c *MongoConnector) PullRecords(
 
 	changeStreamOpts := options.ChangeStream().
 		SetComment("PeerDB changeStream for mirror " + req.FlowJobName).
-		SetFullDocument(options.UpdateLookup)
+		SetFullDocument(options.UpdateLookup).
+		// batchSize=0 only affects the initial aggregate response, so Watch returns
+		// immediately without blocking on the initial cursor establishment. Subsequent
+		// getMore calls fall back to the server default (up to 16 MiB per batch).
+		// https://www.mongodb.com/docs/manual/reference/method/cursor.batchSize/
+		SetBatchSize(0)
 
 	var resumeToken bson.Raw
 	var err error
@@ -157,7 +178,7 @@ func (c *MongoConnector) PullRecords(
 		return err
 	}
 
-	changeStream, err := createChangeStream(c.client, ctx, pipeline, changeStreamOpts)
+	changeStream, err := c.createChangeStream(ctx, pipeline, changeStreamOpts)
 	if err != nil {
 		if isResumeTokenNotFoundError(err) && resumeToken != nil {
 			timestamp, err := decodeTimestampFromResumeToken(resumeToken)
@@ -166,7 +187,7 @@ func (c *MongoConnector) PullRecords(
 			}
 			changeStreamOpts.SetStartAtOperationTime(&timestamp)
 			changeStreamOpts.SetResumeAfter(nil)
-			changeStream, err = createChangeStream(c.client, ctx, pipeline, changeStreamOpts)
+			changeStream, err = c.createChangeStream(ctx, pipeline, changeStreamOpts)
 			if err != nil {
 				return fmt.Errorf("failed to recreate change stream: %w", err)
 			}
@@ -207,12 +228,23 @@ func (c *MongoConnector) PullRecords(
 		otelManager.Metrics.AllFetchedBytesCounter.Add(ctx, read)
 	}()
 
-	checkpoint := func() {
-		if resumeToken := changeStream.ResumeToken(); resumeToken != nil {
-			resumeTokenText := base64.StdEncoding.EncodeToString(resumeToken)
-			req.RecordStream.UpdateLatestCheckpointText(resumeTokenText)
-		} else {
+	checkpoint := func() string {
+		rt := changeStream.ResumeToken()
+		if rt == nil {
 			c.logger.Warn("change stream does not currently contain a resume token")
+			return ""
+		}
+		text := base64.StdEncoding.EncodeToString(rt)
+		req.RecordStream.UpdateLatestCheckpointText(text)
+		return text
+	}
+	checkpointToCatalog := func() {
+		text := checkpoint()
+		if text == "" {
+			return
+		}
+		if err := c.metadataStore.SetLastOffset(ctx, req.FlowJobName, model.CdcCheckpoint{Text: text}); err != nil {
+			c.logger.Error("failed to persist resume token", slog.String("resumeToken", text), slog.Any("error", err))
 		}
 	}
 
@@ -300,7 +332,7 @@ func (c *MongoConnector) PullRecords(
 			changeStreamOpts.SetStartAtOperationTime(nil)
 		}
 
-		changeStream, err = createChangeStream(c.client, ctx, pipeline, changeStreamOpts)
+		changeStream, err = c.createChangeStream(ctx, pipeline, changeStreamOpts)
 		if err != nil {
 			return err
 		}
@@ -317,11 +349,14 @@ func (c *MongoConnector) PullRecords(
 
 			if errors.Is(err, context.DeadlineExceeded) {
 				if recordCount > 0 {
+					// advance offset to the PostBatchResumeToken since the last change event's resume token may be quite old
+					checkpoint()
 					break
 				}
-				// update with PostBatchResumeToken on empty batch
-				// ref: https://github.com/mongodb/specifications/blob/master/source/change-streams/change-streams.md
-				checkpoint()
+				// when no events arrived in this batch, still advance offset to the PostBatchResumeToken.
+				// it's safe to persist to catalog since no records were handed off to the sync workflow,
+				// so there's no in-flight data we could skip past by advancing the offset.
+				checkpointToCatalog()
 				// DeadlineExceeded errors are deemed not recoverable/resumable, so we have to create a new change stream instance
 				if err := recreateChangeStream(false); err != nil {
 					return fmt.Errorf("failed to recreate change stream: %w", err)
@@ -342,12 +377,13 @@ func (c *MongoConnector) PullRecords(
 			return fmt.Errorf("change stream error: %w", err)
 		}
 
-		changeEventSize := int64(len(changeStream.Current))
+		current := changeStream.Current()
+		changeEventSize := int64(len(current))
 		deltaBytesProcessed.Add(changeEventSize)
 		cumulativeBytesProcessed.Add(changeEventSize)
 
 		var changeEvent ChangeEvent
-		if err := bson.Unmarshal(changeStream.Current, &changeEvent); err != nil {
+		if err := bson.Unmarshal(current, &changeEvent); err != nil {
 			return fmt.Errorf("failed to decode change stream document: %w", err)
 		}
 
@@ -465,20 +501,6 @@ func createPipeline(tableNameMapping map[string]model.NameAndExclude) (mongo.Pip
 	return pipeline, nil
 }
 
-// createChangeStream calls client.Watch with a 5 minute context deadline
-// so the driver sends it as maxTimeMS over the wire. Otherwise, server-side
-// default maxTimeMS is used which can sometimes be too short.
-func createChangeStream(
-	client *mongo.Client,
-	parent context.Context,
-	pipeline mongo.Pipeline,
-	opts *options.ChangeStreamOptionsBuilder,
-) (*mongo.ChangeStream, error) {
-	watchCtx, cancel := context.WithTimeout(parent, 5*time.Minute)
-	defer cancel()
-	return client.Watch(watchCtx, pipeline, opts)
-}
-
 // This can happen if the resumeToken we are attempting to `ResumeAfter` refers to a table that has been
 // filtered out of the change stream pipeline (for example, if a user pauses and edits a mirror). If
 // this happens, we decode the resumeToken and extract its operation time, and start a new changeStream

flow/connectors/mongo/cdc_test.go

@@ -0,0 +1,180 @@
+package connmongo
+
+import (
+	"context"
+	"encoding/base64"
+	"encoding/binary"
+	"encoding/hex"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	"go.mongodb.org/mongo-driver/v2/bson"
+	"go.mongodb.org/mongo-driver/v2/mongo"
+	"go.mongodb.org/mongo-driver/v2/mongo/options"
+
+	"github.com/PeerDB-io/peerdb/flow/generated/protos"
+	"github.com/PeerDB-io/peerdb/flow/internal"
+	"github.com/PeerDB-io/peerdb/flow/model"
+	"github.com/PeerDB-io/peerdb/flow/otel_metrics"
+	"github.com/PeerDB-io/peerdb/flow/shared"
+)
+
+type iterationType int
+
+const (
+	// idle returns false on Next() with context.DeadlineExceeded
+	idle iterationType = iota
+	// insert returns true on Next() with a generated insert event
+	insert
+)
+
+//nolint:govet // it's a test, no need for fieldalignment
+type mockChangeStream struct {
+	err         error
+	resumeToken bson.Raw
+	current     bson.Raw
+
+	idx          int
+	iterations   []iterationType
+	emittedTimes []time.Time
+
+	t *testing.T
+}
+
+func newMockChangeStream(t *testing.T, iter ...iterationType) *mockChangeStream {
+	t.Helper()
+	return &mockChangeStream{t: t, iterations: iter}
+}
+
+func (cs *mockChangeStream) Next(context.Context) bool {
+	if cs.idx >= len(cs.iterations) {
+		cs.t.Fatalf("mockChangeStream: Next past end of mocked iterations (%d iterations)", len(cs.iterations))
+	}
+	ts := time.Now()
+	cs.emittedTimes = append(cs.emittedTimes, ts)
+	cs.resumeToken = toResumeToken(ts)
+
+	label := cs.iterations[cs.idx]
+	cs.idx++
+
+	switch label {
+	case insert:
+		cs.current = newInsertChangeEvent(bson.NewObjectID(), ts)
+		cs.err = nil
+		return true
+	case idle:
+		cs.err = context.DeadlineExceeded
+		return false
+	default:
+		cs.t.Fatalf("mockChangeStream: unknown label %d", label)
+		return false
+	}
+}
+
+func (cs *mockChangeStream) ResumeToken() bson.Raw       { return cs.resumeToken }
+func (cs *mockChangeStream) Err() error                  { return cs.err }
+func (cs *mockChangeStream) Current() bson.Raw           { return cs.current }
+func (cs *mockChangeStream) Close(context.Context) error { return nil }
+
+var _ ChangeStream = (*mockChangeStream)(nil)
+
+type mockMetadataStore struct{ persisted []model.CdcCheckpoint }
+
+func (ms *mockMetadataStore) GetLastOffset(context.Context, string) (model.CdcCheckpoint, error) {
+	if n := len(ms.persisted); n > 0 {
+		return ms.persisted[n-1], nil
+	}
+	return model.CdcCheckpoint{}, nil
+}
+
+func (ms *mockMetadataStore) SetLastOffset(_ context.Context, _ string, off model.CdcCheckpoint) error {
+	ms.persisted = append(ms.persisted, off)
+	return nil
+}
+
+func drainMongoCDCRecordsAsync(t *testing.T, stream *model.CDCStream[model.RecordItems]) {
+	t.Helper()
+	go func() {
+		for range stream.GetRecords() {
+		}
+	}()
+}
+
+func newInsertChangeEvent(id bson.ObjectID, ts time.Time) bson.Raw {
+	event, _ := bson.Marshal(bson.D{
+		{Key: "ns", Value: bson.D{
+			{Key: "db", Value: "db"},
+			{Key: "coll", Value: "coll"},
+		}},
+		{Key: "operationType", Value: "insert"},
+		{Key: "documentKey", Value: bson.D{{Key: "_id", Value: id}}},
+		{Key: "fullDocument", Value: bson.D{
+			{Key: "_id", Value: id},
+			{Key: "val", Value: "test"},
+		}},
+		{Key: "clusterTime", Value: toBsonTs(ts)},
+	})
+	return event
+}
+
+func TestChangeStreamIdleConnectionAdvancesOffset(t *testing.T) {
+	ctx := t.Context()
+
+	mockCS := newMockChangeStream(t, idle, idle, insert, idle)
+	mockStore := &mockMetadataStore{}
+	connector := &MongoConnector{
+		logger: internal.LoggerFromCtx(t.Context()),
+		createChangeStream: func(
+			context.Context, mongo.Pipeline, ...options.Lister[options.ChangeStreamOptions],
+		) (ChangeStream, error) {
+			return mockCS, nil
+		},
+		metadataStore: mockStore,
+	}
+
+	otelManager, err := otel_metrics.NewOtelManager(ctx, "test", false)
+	require.NoError(t, err)
+
+	req := &model.PullRecordsRequest[model.RecordItems]{
+		FlowJobName:            "test_mongo_idle",
+		RecordStream:           model.NewCDCStream[model.RecordItems](100),
+		TableNameMapping:       map[string]model.NameAndExclude{"db.coll": {Name: "db_coll"}},
+		TableNameSchemaMapping: map[string]*protos.TableSchema{},
+		MaxBatchSize:           10000,
+		IdleTimeout:            time.Minute,
+		Env:                    map[string]string{"PEERDB_MONGODB_DIRECT_BSON_CONVERTER": "true"},
+	}
+	drainMongoCDCRecordsAsync(t, req.RecordStream)
+
+	require.NoError(t, connector.PullRecords(ctx, shared.CatalogPool{}, otelManager, req))
+	require.Len(t, mockStore.persisted, 2)
+	require.Equal(t, b64(toResumeToken(mockCS.emittedTimes[0])), mockStore.persisted[0].Text)
+	require.Equal(t, b64(toResumeToken(mockCS.emittedTimes[1])), mockStore.persisted[1].Text)
+	require.Equal(t, b64(toResumeToken(mockCS.emittedTimes[3])), req.RecordStream.GetLastCheckpoint().Text)
+}
+
+func b64(raw bson.Raw) string {
+	return base64.StdEncoding.EncodeToString(raw)
+}
+
+func toBsonTs(ts time.Time) bson.Timestamp {
+	return bson.Timestamp{T: uint32(ts.Unix()), I: uint32(ts.Nanosecond())}
+}
+
+func toResumeToken(ts time.Time) bson.Raw {
+	t := toBsonTs(ts)
+	keyString := make([]byte, 9)
+	keyString[0] = byte(kTimestamp)
+	binary.BigEndian.PutUint64(keyString[1:], uint64(t.T)<<32|uint64(t.I))
+	raw, _ := bson.Marshal(bson.D{{Key: "_data", Value: hex.EncodeToString(keyString)}})
+	return raw
+}
+
+func TestResumeTokenHelpersRoundTrip(t *testing.T) {
+	ts := time.Now().UTC()
+	rt := toResumeToken(ts)
+	bsonTs, err := decodeTimestampFromResumeToken(rt)
+	require.NoError(t, err)
+	require.Equal(t, toBsonTs(ts), bsonTs)
+}