diff --git a/flow/connectors/mongo/cdc.go b/flow/connectors/mongo/cdc.go index 203ab911c7..6f464a6816 100644 --- a/flow/connectors/mongo/cdc.go +++ b/flow/connectors/mongo/cdc.go @@ -24,6 +24,7 @@ import ( "github.com/PeerDB-io/peerdb/flow/otel_metrics" "github.com/PeerDB-io/peerdb/flow/pkg/common" "github.com/PeerDB-io/peerdb/flow/shared" + "github.com/PeerDB-io/peerdb/flow/shared/concurrency" "github.com/PeerDB-io/peerdb/flow/shared/exceptions" "github.com/PeerDB-io/peerdb/flow/shared/types" ) @@ -230,6 +231,132 @@ func decodeEvent( return nil } +// Constant used by PullRecords. +// +// Number of recordItems to pass in one chunk to decode/send workers managed by PullRecordsWorkerPool. +// Doing per-item channel sends results in too much coordination and reduces effective concurrency +// in practice. +const pullRecordsItemsChunkSize = 256 + +// PullRecords spins up worker goroutines using PullRecordsWorkerPool. +// The context passed into these goroutines is not used to signal timeouts; +// rather, PullRecordsWorkerPool ensures that after Wait() is called on it, +// we gracefully drain all items through this function. +func (c *MongoConnector) recordSender( + ctx context.Context, + records []model.Record[model.RecordItems], + resumeToken string, + req *model.PullRecordsRequest[model.RecordItems], + signalledAsNonEmpty *bool, +) error { + for i := range records { + if !*signalledAsNonEmpty { + // This bool should be shared across any instantiations of sendLoop for + // a given RecordStream. However, it's not an atomic and so only one sendLoop + // at a given time can own it. + *signalledAsNonEmpty = true + req.RecordStream.SignalAsNotEmpty() + } + if err := req.RecordStream.AddRecord(ctx, records[i]); err != nil { + return err + } + } + if resumeToken != "" { + req.RecordStream.UpdateLatestCheckpointText(resumeToken) + } + return nil +} + +type encodedMongoEvent struct { + maybeFullDocument *bson.Raw + operationType operationType + sourceTableName string + destinationTableName string + documentKey bson.Raw + commitTimeNanos int64 +} + +// decodeEvent is spun up by PullRecordsWorkerPool in separate goroutines, up to +// PEERDB_MONGODB_NUM_PARALLEL_DECODE_THREADS in parallel. The output is sent to `recordSender` +// in order. +func (c *MongoConnector) decodeEvent( + events []encodedMongoEvent, + req *model.PullRecordsRequest[model.RecordItems], +) ([]model.Record[model.RecordItems], error) { + // Utils used by this routine. + converter := NewDirectBsonConverter() + fullDocumentColumnName := DefaultFullDocumentColumnName + if req.InternalVersion < shared.InternalVersion_MongoDBFullDocumentColumnToDoc { + fullDocumentColumnName = LegacyFullDocumentColumnName + } + parseItem := func(event encodedMongoEvent) (model.Record[model.RecordItems], error) { + items := model.NewRecordItems(2) + + if len(event.documentKey) > 0 { + rv := event.documentKey.Lookup(DefaultDocumentKeyColumnName) + if rv.IsZero() || rv.Type == bson.TypeNull { + return nil, exceptions.NewInvalidIdValueError(event.sourceTableName) + } + qValue, err := converter.QValueStringFromId(rv, req.InternalVersion) + if err != nil { + return nil, fmt.Errorf("failed to convert key: %w", err) + } + items.AddColumn(DefaultDocumentKeyColumnName, qValue) + } else { + return nil, fmt.Errorf("document key is nil") + } + + if event.maybeFullDocument != nil && len(*event.maybeFullDocument) > 0 { + qValue, err := converter.QValueJSONFromDocument(*event.maybeFullDocument) + if err != nil { + return nil, fmt.Errorf("failed to convert document: %w", err) + } + items.AddColumn(fullDocumentColumnName, qValue) + } else { + // `fullDocument` field will not exist in the following scenarios: + // 1) operationType is 'delete' + // 2) document is deleted / collection is dropped in between update and lookup + // 3) update changes the values for at least one of the fields in that collection's + // shard key (although sharding is not supported today) + items.AddColumn(fullDocumentColumnName, types.QValueJSON{Val: "{}"}) + } + var record model.Record[model.RecordItems] + switch event.operationType { + case operationTypeInsert: + record = &model.InsertRecord[model.RecordItems]{ + BaseRecord: model.BaseRecord{CommitTimeNano: event.commitTimeNanos}, + Items: items, + SourceTableName: event.sourceTableName, + DestinationTableName: event.destinationTableName, + } + + case operationTypeUpdate, operationTypeReplace: + record = &model.UpdateRecord[model.RecordItems]{ + BaseRecord: model.BaseRecord{CommitTimeNano: event.commitTimeNanos}, + NewItems: items, + SourceTableName: event.sourceTableName, + DestinationTableName: event.destinationTableName, + } + case operationTypeDelete: + record = &model.DeleteRecord[model.RecordItems]{ + BaseRecord: model.BaseRecord{CommitTimeNano: event.commitTimeNanos}, + Items: items, + SourceTableName: event.sourceTableName, + DestinationTableName: event.destinationTableName, + } + } + return record, nil + } + modelRecords := make([]model.Record[model.RecordItems], len(events)) + for i := range events { + var err error + if modelRecords[i], err = parseItem(events[i]); err != nil { + return nil, err + } + } + return modelRecords, nil +} + func (c *MongoConnector) PullRecords( ctx context.Context, catalogPool shared.CatalogPool, @@ -243,11 +370,6 @@ func (c *MongoConnector) PullRecords( alerter = alerting.NewAlerter(ctx, catalogPool, otelManager) } - fullDocumentColumnName := DefaultFullDocumentColumnName - if req.InternalVersion < shared.InternalVersion_MongoDBFullDocumentColumnToDoc { - fullDocumentColumnName = LegacyFullDocumentColumnName - } - c.logger.Info("[mongo] started PullRecords for mirror "+req.FlowJobName, slog.Any("table_mapping", req.TableNameMapping), slog.Uint64("max_batch_size", uint64(req.MaxBatchSize)), @@ -314,6 +436,7 @@ func (c *MongoConnector) PullRecords( var recordCount uint32 var deltaBytesProcessed, cumulativeBytesProcessed atomic.Int64 + var signalledAsNonEmpty bool pullStart := time.Now() defer func() { if recordCount == 0 { @@ -339,6 +462,21 @@ func (c *MongoConnector) PullRecords( slog.Int("channelLen", req.RecordStream.ChannelLen()), slog.Float64("elapsedMinutes", time.Since(pullStart).Minutes())) }() + numParallelDecodeWorkers, err := internal.PeerDBMongoDBNumParallelDecodeThreads(ctx, req.Env) + if err != nil { + return err + } + workerPool := concurrency.PullRecordsWorkerPool[encodedMongoEvent, []model.Record[model.RecordItems], string]{ + Concurrency: int(numParallelDecodeWorkers), + ChunkSize: pullRecordsItemsChunkSize, + WorkerFunc: func(events []encodedMongoEvent) ([]model.Record[model.RecordItems], error) { + return c.decodeEvent(events, req) + }, + Send: func(ctx context.Context, items []model.Record[model.RecordItems], resumeToken string) error { + return c.recordSender(ctx, items, resumeToken, req, &signalledAsNonEmpty) + }, + } + workerPool.Init(ctx) // before the first record arrives, we wait for up to an hour before resetting context timeout // after the first record arrives, we switch to configured idleTimeout timeoutCtx, cancelTimeout := context.WithTimeout(ctx, time.Hour) @@ -351,6 +489,7 @@ func (c *MongoConnector) PullRecords( defer func() { cancelTimeout() + _ = workerPool.Wait(ctx) reportBytesShutdown() read := deltaBytesProcessed.Swap(0) otelManager.Metrics.FetchedBytesCounter.Add(ctx, read) @@ -377,46 +516,9 @@ func (c *MongoConnector) PullRecords( } } - converter := NewDirectBsonConverter() - addRecordItems := func(documentKey bson.Raw, maybeFullDocument *bson.Raw, items *model.RecordItems, tableName string) error { - if len(documentKey) > 0 { - rv := documentKey.Lookup(DefaultDocumentKeyColumnName) - if rv.IsZero() || rv.Type == bson.TypeNull { - return exceptions.NewInvalidIdValueError(tableName) - } - qValue, err := converter.QValueStringFromId(rv, req.InternalVersion) - if err != nil { - return fmt.Errorf("failed to convert key: %w", err) - } - items.AddColumn(DefaultDocumentKeyColumnName, qValue) - } else { - return fmt.Errorf("document key is nil") - } - - if maybeFullDocument != nil && len(*maybeFullDocument) > 0 { - qValue, err := converter.QValueJSONFromDocument(*maybeFullDocument) - if err != nil { - return fmt.Errorf("failed to convert document: %w", err) - } - items.AddColumn(fullDocumentColumnName, qValue) - } else { - // `fullDocument` field will not exist in the following scenarios: - // 1) operationType is 'delete' - // 2) document is deleted / collection is dropped in between update and lookup - // 3) update changes the values for at least one of the fields in that collection's - // shard key (although sharding is not supported today) - items.AddColumn(fullDocumentColumnName, types.QValueJSON{Val: "{}"}) - } - return nil - } - - addRecord := func(ctx context.Context, record model.Record[model.RecordItems]) error { + incrementRecordCount := func() { recordCount += 1 - if err := req.RecordStream.AddRecord(ctx, record); err != nil { - return err - } if recordCount == 1 { - req.RecordStream.SignalAsNotEmpty() timeoutCtx, cancelTimeout = context.WithTimeout(ctx, req.IdleTimeout) //nolint:gosec // G118: cancelTimeout called in defer } if recordCount%50000 == 0 { @@ -426,7 +528,6 @@ func (c *MongoConnector) PullRecords( slog.Int("channelLen", req.RecordStream.ChannelLen()), slog.Float64("elapsedMinutes", time.Since(pullStart).Minutes())) } - return nil } recreateChangeStream := func(useOperationTime bool) error { @@ -445,6 +546,9 @@ func (c *MongoConnector) PullRecords( cancelTimeout() timeoutCtx, cancelTimeout = context.WithTimeout(ctx, time.Hour) + // reset worker pool. Wait() has already been called on this workerPool. + workerPool.Init(ctx) + // set resume point based on whether operation time should be used or not if useOperationTime { timestamp, err := decodeTimestampFromResumeToken(resumeToken) @@ -474,11 +578,22 @@ func (c *MongoConnector) PullRecords( return fmt.Errorf("unexpected: changestream.Next() returned false but no change stream error was recorded") } + if err := workerPool.Flush(ctx); err != nil { + return err + } + + if err := workerPool.Wait(ctx); err != nil { + return err + } + 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 + // + // This checkpoint is safe to do here as opposed to in sendLoop, because sendLoop has been drained away + // above. checkpoint() - break + return nil } // 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, @@ -539,47 +654,18 @@ func (c *MongoConnector) PullRecords( continue } - items := model.NewRecordItems(2) - switch operationType(changeEvent.OperationType) { - case operationTypeInsert: - if err := addRecordItems(changeEvent.DocumentKey, changeEvent.FullDocument, &items, sourceTableName); err != nil { - return fmt.Errorf("failed to process document: %w", err) - } - - if err = addRecord(ctx, &model.InsertRecord[model.RecordItems]{ - BaseRecord: model.BaseRecord{CommitTimeNano: commitTimeNanos}, - Items: items, - SourceTableName: sourceTableName, - DestinationTableName: destinationTableName, - }); err != nil { - return fmt.Errorf("failed to add insert record: %w", err) - } - case operationTypeUpdate, operationTypeReplace: - if err := addRecordItems(changeEvent.DocumentKey, changeEvent.FullDocument, &items, sourceTableName); err != nil { - return fmt.Errorf("failed to process document: %w", err) - } - - if err := addRecord(ctx, &model.UpdateRecord[model.RecordItems]{ - BaseRecord: model.BaseRecord{CommitTimeNano: commitTimeNanos}, - NewItems: items, - SourceTableName: sourceTableName, - DestinationTableName: destinationTableName, - }); err != nil { - return fmt.Errorf("failed to add update record: %w", err) - } - case operationTypeDelete: - if err := addRecordItems(changeEvent.DocumentKey, changeEvent.FullDocument, &items, sourceTableName); err != nil { - return fmt.Errorf("failed to process document: %w", err) - } - - if err := addRecord(ctx, &model.DeleteRecord[model.RecordItems]{ - BaseRecord: model.BaseRecord{CommitTimeNano: commitTimeNanos}, - Items: items, - SourceTableName: sourceTableName, - DestinationTableName: destinationTableName, - }); err != nil { - return fmt.Errorf("failed to add delete record: %w", err) - } + event := encodedMongoEvent{ + documentKey: changeEvent.DocumentKey, + maybeFullDocument: changeEvent.FullDocument, + operationType: operationType(changeEvent.OperationType), + sourceTableName: sourceTableName, + destinationTableName: destinationTableName, + commitTimeNanos: commitTimeNanos, + } + switch event.operationType { + case operationTypeInsert, operationTypeReplace, operationTypeUpdate, operationTypeDelete: + // Happy path. + incrementRecordCount() default: c.logger.Warn(fmt.Sprintf("skipping event with unsupported operation type '%s' (db=%s coll=%s)", changeEvent.OperationType, changeEvent.Ns.Db, changeEvent.Ns.Coll)) @@ -600,7 +686,23 @@ func (c *MongoConnector) PullRecords( continue } otelManager.Metrics.FetchedEventSizeHistogram.Record(ctx, changeEventSize) - checkpoint() + rt := changeStream.ResumeToken() + var rtText string + if rt == nil { + c.logger.Warn("change stream does not currently contain a resume token") + } else { + rtText = base64.StdEncoding.EncodeToString(rt) + } + if err := workerPool.AddItem(ctx, event, rtText); err != nil { + return err + } + } + if err := workerPool.Flush(ctx); err != nil { + return err + } + + if err := workerPool.Wait(ctx); err != nil { + return err } return nil diff --git a/flow/connectors/mongo/cdc_batch_test.go b/flow/connectors/mongo/cdc_batch_test.go new file mode 100644 index 0000000000..6227c95157 --- /dev/null +++ b/flow/connectors/mongo/cdc_batch_test.go @@ -0,0 +1,373 @@ +package connmongo + +import ( + "context" + "fmt" + "slices" + "testing" + "time" + + "github.com/stretchr/testify/require" + "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" +) + +// pullHarness wires a MongoConnector up to a scripted change stream and an in-memory +// metadata store, so a whole PullRecords call can be driven without a server. +type pullHarness struct { + connector *MongoConnector + store *mockMetadataStore + stream *mockChangeStream + req *model.PullRecordsRequest[model.RecordItems] + // cancel cancels the context PullRecords runs under, so a test can cancel the + // activity mid-pull. It is set by run, and so is only safe to call from inside a + // pull. + cancel context.CancelFunc + // streamCreations counts createChangeStream calls: one for the initial stream plus + // one per recreation. + streamCreations int +} + +// pullOutcome is everything the sync workflow gets to observe from one pull: the +// records handed to the record stream, the offset left on it, and the error, if any. +type pullOutcome struct { + err error + checkpoint string + ids []string +} + +func newPullHarness(t *testing.T, iterations ...iterationType) *pullHarness { + t.Helper() + h := &pullHarness{ + store: &mockMetadataStore{}, + stream: newMockChangeStream(t, iterations...), + } + h.connector = &MongoConnector{ + logger: internal.LoggerFromCtx(t.Context()), + createChangeStream: func( + context.Context, mongo.Pipeline, ...options.Lister[options.ChangeStreamOptions], + ) (ChangeStream, error) { + h.streamCreations++ + return h.stream, nil + }, + metadataStore: h.store, + } + h.req = &model.PullRecordsRequest[model.RecordItems]{ + FlowJobName: t.Name(), + // Buffered past any batch these tests pull, so that a blocked AddRecord means + // a real deadlock rather than a slow drain goroutine. + RecordStream: model.NewCDCStream[model.RecordItems](4096), + TableNameMapping: map[string]model.NameAndExclude{"db.coll": {Name: "db_coll"}}, + TableNameSchemaMapping: map[string]*protos.TableSchema{}, + MaxBatchSize: 10000, + IdleTimeout: time.Minute, + InternalVersion: shared.InternalVersion_Latest, + } + return h +} + +// run drives PullRecords to completion, draining the record stream concurrently the way +// the sync side of the activity does. +func (h *pullHarness) run(t *testing.T) pullOutcome { + t.Helper() + + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + h.cancel = cancel + + ids := []string{} + drained := make(chan struct{}) + go func() { + defer close(drained) + for record := range h.req.RecordStream.GetRecords() { + if record == nil { + ids = append(ids, "") + continue + } + value, err := record.GetItems().GetValueByColName(DefaultDocumentKeyColumnName) + if err != nil { + ids = append(ids, fmt.Sprintf("", err)) + continue + } + id, ok := value.Value().(string) + if !ok { + ids = append(ids, fmt.Sprintf("", value.Value())) + continue + } + ids = append(ids, id) + } + }() + + otelManager, err := otel_metrics.NewOtelManager(t.Context(), "test", false) + require.NoError(t, err) + pullErr := h.connector.PullRecords(ctx, shared.CatalogPool{}, otelManager, h.req) + <-drained + + return pullOutcome{err: pullErr, checkpoint: h.req.RecordStream.GetLastCheckpoint().Text, ids: ids} +} + +// persistedOffsets is every offset the pull wrote straight to the catalog. +func (h *pullHarness) persistedOffsets() []string { + offsets := make([]string, 0, len(h.store.persisted)) + for _, persisted := range h.store.persisted { + offsets = append(offsets, persisted.Text) + } + return offsets +} + +// requireOffsetsCoverOnlyDeliveredRecords asserts the safety property PullRecords owes +// the sync workflow: a resume token once every replicable event up to +// it has been handed to the RecordStream. The next batch resumes strictly after the +// committed offset, so an offset that runs ahead of the delivered records skips those +// events for good. +// +// Two kinds of offset are checked, and only these two are ever committed: +// - offsets PullRecords writes straight to the catalog, which are committed the moment +// they are written and so are never exempt; +// - the offset left on the record stream by a pull that returned no error, which the +// sync workflow commits once the destination has taken the batch. +func (h *pullHarness) requireOffsetsCoverOnlyDeliveredRecords(t *testing.T, out pullOutcome) { + t.Helper() + + offsets := h.persistedOffsets() + if out.err == nil && out.checkpoint != "" { + offsets = append(offsets, out.checkpoint) + } + + for _, offset := range offsets { + want := h.stream.replicableThrough(h.stream.iterationOfToken(offset)) + require.LessOrEqual(t, len(want), len(out.ids), + "offset covers %d replicable events but only %d records reached the record stream", len(want), len(out.ids)) + require.Equal(t, want, out.ids[:len(want)], + "offset commits past events that never reached the record stream") + } +} + +// repeatInserts scripts n insert iterations. +func repeatInserts(n int) []iterationType { + return slices.Repeat([]iterationType{insert}, n) +} + +// mockEventIDs is the _id of the first n insert events. +func mockEventIDs(n int) []string { + ids := make([]string, 0, n) + for i := range n { + ids = append(ids, mockEventID(i)) + } + return ids +} + +func TestPullRecordsOffsetNeverRunsAheadOfDeliveredRecords(t *testing.T) { + for _, tc := range []struct { + name string + iterations []iterationType + maxBatchSize uint32 + }{ + {name: "cut by max batch size", iterations: repeatInserts(8), maxBatchSize: 5}, + { + name: "cut by max batch size across decode batches", + iterations: repeatInserts(pullRecordsItemsChunkSize + 9), + maxBatchSize: pullRecordsItemsChunkSize + 3, + }, + {name: "cut by sync interval", iterations: []iterationType{insert, insert, idle}, maxBatchSize: 100}, + {name: "idle before any record", iterations: []iterationType{idle, idle, insert, idle}, maxBatchSize: 100}, + {name: "idle after every record", iterations: []iterationType{insert, idle}, maxBatchSize: 100}, + {name: "fatal error mid batch", iterations: []iterationType{insert, insert, fatal}, maxBatchSize: 100}, + {name: "fatal error before any record", iterations: []iterationType{fatal}, maxBatchSize: 100}, + {name: "undecodable document", iterations: []iterationType{insert, insert, nullIdInsert}, maxBatchSize: 3}, + { + name: "undecodable document after a full decode batch", + iterations: append(repeatInserts(pullRecordsItemsChunkSize), nullIdInsert), + maxBatchSize: pullRecordsItemsChunkSize + 1, + }, + // A batch that never sees a record keeps recreating its stream on every idle + // timeout, so these cases have to be ended by a hard failure instead. + {name: "unsupported operations only", iterations: []iterationType{unsupportedOp, unsupportedOp, fatal}, maxBatchSize: 100}, + } { + t.Run(tc.name, func(t *testing.T) { + h := newPullHarness(t, tc.iterations...) + h.req.MaxBatchSize = tc.maxBatchSize + + out := h.run(t) + // Whatever reached the stream must be a gap-free prefix of the emitted + // inserts: a hole would mean an offset skipped an event even if the + // offset itself looks sane. + require.Equal(t, mockEventIDs(len(out.ids)), out.ids, "records reached the stream out of order or with gaps") + h.requireOffsetsCoverOnlyDeliveredRecords(t, out) + }) + } +} + +// MaxBatchSize is a hard cap: PullRecords must stop pulling the moment it is reached, +// because an event read past the cap is neither delivered nor covered by the offset, and +// the change stream it came from is closed on the way out. +func TestPullRecordsTruncatesBatchAtMaxBatchSize(t *testing.T) { + const maxBatchSize = 4 + h := newPullHarness(t, repeatInserts(maxBatchSize+6)...) + h.req.MaxBatchSize = maxBatchSize + + out := h.run(t) + require.NoError(t, out.err) + require.Equal(t, mockEventIDs(maxBatchSize), out.ids) + require.Equal(t, maxBatchSize, h.stream.idx, "pulled events past MaxBatchSize") + require.Equal(t, h.stream.tokenAt(maxBatchSize-1), out.checkpoint) + require.Empty(t, h.persistedOffsets(), "a full batch is committed by the sync workflow, not the puller") + h.requireOffsetsCoverOnlyDeliveredRecords(t, out) +} + +// A batch larger than pullRecordsItemsChunkSize is handed to the decode workers in +// several chunks that are decoded in parallel. Every event must still arrive exactly +// once and in change stream order, and the offset must land on the last one. +func TestPullRecordsTruncatesBatchSpanningManyDecodeBatches(t *testing.T) { + maxBatchSize := 2*pullRecordsItemsChunkSize + 37 + h := newPullHarness(t, repeatInserts(maxBatchSize+1)...) + h.req.MaxBatchSize = uint32(maxBatchSize) + + out := h.run(t) + require.NoError(t, out.err) + require.Equal(t, mockEventIDs(maxBatchSize), out.ids) + require.Equal(t, maxBatchSize, h.stream.idx, "pulled events past MaxBatchSize") + require.Equal(t, h.stream.tokenAt(maxBatchSize-1), out.checkpoint) + h.requireOffsetsCoverOnlyDeliveredRecords(t, out) +} + +// The sync interval bounds the batch from the arrival of its first record, not the gap +// between consecutive records: once the interval elapses the batch is cut and handed +// over, however busy the stream still is. +func TestPullRecordsSyncIntervalBoundsBatchFromFirstRecord(t *testing.T) { + const syncInterval = 90 * time.Second + h := newPullHarness(t, insert, insert, idle) + h.req.IdleTimeout = syncInterval + + out := h.run(t) + require.NoError(t, out.err) + require.Equal(t, mockEventIDs(2), out.ids) + + // Before any record arrives the pull waits on the long, hour-long budget rather + // than the sync interval, so an idle mirror does not spin recreating its stream. + require.Greater(t, h.stream.deadlines[0], time.Hour-time.Minute) + // Once a record has landed, the remaining budget is the configured sync interval, + // and it keeps counting down rather than being reset per record. + require.LessOrEqual(t, h.stream.deadlines[1], syncInterval) + require.Greater(t, h.stream.deadlines[1], syncInterval-time.Minute) + require.Less(t, h.stream.deadlines[2], h.stream.deadlines[1], "sync interval restarted mid-batch") + + // Cutting a non-empty batch advances the offset to the post-batch resume token, + // which is safe because every record is already on the stream, but it must be left + // for the sync workflow to commit rather than persisted here. + require.Equal(t, h.stream.tokenAt(2), out.checkpoint) + require.Empty(t, h.persistedOffsets()) + require.Equal(t, 1, h.streamCreations, "cutting a batch should not recreate the change stream") + h.requireOffsetsCoverOnlyDeliveredRecords(t, out) +} + +// An idle stream with nothing pulled yet is the one case where the puller commits an +// offset itself. That is safe because no records are in flight, and it is what +// stops an idle mirror's resume token from ageing out. A batch that has seen no record +// keeps waiting rather than returning empty, so this one only ends once a record shows up. +func TestPullRecordsIdleWithEmptyBatchPersistsOffsetAndRecreatesStream(t *testing.T) { + h := newPullHarness(t, idle, idle, insert, idle) + + out := h.run(t) + require.NoError(t, out.err) + require.Equal(t, mockEventIDs(1), out.ids) + // One offset per idle timeout that found the batch empty, and none for the timeout + // that cut the batch holding a record: committing that one is the sync workflow's job. + require.Equal(t, []string{h.stream.tokenAt(0), h.stream.tokenAt(1)}, h.persistedOffsets()) + require.Equal(t, h.stream.tokenAt(3), out.checkpoint) + // One initial stream plus one recreation per idle timeout: a deadline-exceeded + // stream is not resumable. + require.Equal(t, 3, h.streamCreations) + h.requireOffsetsCoverOnlyDeliveredRecords(t, out) +} + +// Events the connector cannot replicate must not consume batch budget, and must not +// advance the offset on their own: the offset stays on the last event that was replicated. +func TestPullRecordsSkipsUnsupportedOperations(t *testing.T) { + h := newPullHarness(t, insert, unsupportedOp, unsupportedOp, insert, idle) + h.req.MaxBatchSize = 2 + + out := h.run(t) + require.NoError(t, out.err) + require.Equal(t, mockEventIDs(2), out.ids) + // Five iterations were scripted; the two inserts plus the two skipped events is + // where MaxBatchSize is reached, so the trailing idle is never reached. + require.Equal(t, 4, h.stream.idx) + require.Equal(t, h.stream.tokenAt(3), out.checkpoint) + h.requireOffsetsCoverOnlyDeliveredRecords(t, out) +} + +// A change stream error that is neither a timeout nor a stale resume token ends the +// pull. The caller throws the batch away, so the one thing that must not happen is an +// offset reaching the catalog. +func TestPullRecordsFatalStreamErrorCommitsNoOffset(t *testing.T) { + h := newPullHarness(t, insert, insert, fatal) + + out := h.run(t) + require.ErrorIs(t, out.err, errFatalChangeStream) + require.ErrorContains(t, out.err, "change stream error") + // The workers are drained before the error is surfaced, so the records pulled + // before the failure are intact on the stream rather than half-written. + require.Equal(t, mockEventIDs(2), out.ids) + require.Empty(t, h.persistedOffsets(), "a failed pull must not commit an offset") + h.requireOffsetsCoverOnlyDeliveredRecords(t, out) +} + +// A document the decode workers reject must fail the pull. Silently dropping it would +// leave the offset covering an event that never reached the destination. +func TestPullRecordsUndecodableDocumentFailsPull(t *testing.T) { + h := newPullHarness(t, insert, insert, nullIdInsert) + h.req.MaxBatchSize = 3 + + out := h.run(t) + require.ErrorContains(t, out.err, "_id field is missing or null in table db.coll") + // The bad document is in the only sub-batch, so the whole batch dies with it. + require.Empty(t, out.ids) + require.Empty(t, h.persistedOffsets(), "a failed pull must not commit an offset") + h.requireOffsetsCoverOnlyDeliveredRecords(t, out) +} + +// The same failure in a sub-batch the pull loop has already moved past has to be noticed +// too: the loop is still pulling events when a worker dies behind it, and finishing the +// batch as if nothing happened would hand the sync workflow an offset spanning the +// records that worker dropped. +func TestPullRecordsUndecodableDocumentInEarlierSubBatchFailsPull(t *testing.T) { + // The bad document heads the first sub-batch; MaxBatchSize is two sub-batches, so + // the loop keeps pulling well past the point the worker rejects it. + iterations := append([]iterationType{nullIdInsert}, repeatInserts(3*pullRecordsItemsChunkSize)...) + h := newPullHarness(t, iterations...) + h.req.MaxBatchSize = uint32(2 * pullRecordsItemsChunkSize) + + out := h.run(t) + require.ErrorContains(t, out.err, "_id field is missing or null in table db.coll") + require.Empty(t, h.persistedOffsets(), "a failed pull must not commit an offset") + h.requireOffsetsCoverOnlyDeliveredRecords(t, out) +} + +// Cancelling the activity mid-pull must surface as an error. Returning nil would let the +// sync workflow commit the offset for a batch that was cut off part-way through. +func TestPullRecordsContextCancellationMidBatchFailsPull(t *testing.T) { + h := newPullHarness(t, repeatInserts(4*pullRecordsItemsChunkSize)...) + h.req.MaxBatchSize = 3 * pullRecordsItemsChunkSize + // Cancel once a couple of sub-batches are in flight, so the cancellation races the + // decode workers rather than arriving before any work started. + h.stream.beforeNext = func(idx int) { + if idx == 2*pullRecordsItemsChunkSize { + h.cancel() + } + } + + out := h.run(t) + require.ErrorIs(t, out.err, context.Canceled) + require.Empty(t, h.persistedOffsets(), "a cancelled pull must not commit an offset") + // Whatever did reach the stream is still a gap-free prefix; cancellation drops + // in-flight records rather than reordering or interleaving them. + require.Equal(t, mockEventIDs(len(out.ids)), out.ids) + h.requireOffsetsCoverOnlyDeliveredRecords(t, out) +} diff --git a/flow/connectors/mongo/cdc_test.go b/flow/connectors/mongo/cdc_test.go index 51e7a0cfaa..797f75b265 100644 --- a/flow/connectors/mongo/cdc_test.go +++ b/flow/connectors/mongo/cdc_test.go @@ -6,6 +6,7 @@ import ( "encoding/binary" "encoding/hex" "errors" + "fmt" "log/slog" "strings" "testing" @@ -33,6 +34,15 @@ const ( insert // rawEvent returns true on Next() with the next event queued in mockChangeStream.rawEvents rawEvent + // nullIdInsert returns true on Next() with an insert event whose documentKey._id + // is null, which a decode worker rejects. + nullIdInsert + // unsupportedOp returns true on Next() with an event carrying an operationType + // PullRecords does not replicate. + unsupportedOp + // fatal returns false on Next() with an error that is neither a timeout nor a + // missing resume token, so PullRecords has to give up. + fatal ) type mockChangeStream struct { @@ -44,6 +54,11 @@ type mockChangeStream struct { iterations []iterationType emittedTimes []time.Time rawEvents []bson.Raw + deadlines []time.Duration + inserts int + // beforeNext, when set, runs at the start of each Next() call with the index of the + // iteration about to be served, letting a test interfere mid-pull. + beforeNext func(idx int) t *testing.T } @@ -53,10 +68,25 @@ func newMockChangeStream(t *testing.T, iter ...iterationType) *mockChangeStream return &mockChangeStream{t: t, iterations: iter} } -func (cs *mockChangeStream) Next(context.Context) bool { +func (cs *mockChangeStream) Next(ctx context.Context) bool { + if cs.beforeNext != nil { + cs.beforeNext(cs.idx) + } + // A real cursor cannot serve an event once its context is gone, whatever iterations + // the test still had scripted. + if err := ctx.Err(); err != nil { + cs.err = err + return false + } if cs.idx >= len(cs.iterations) { cs.t.Fatalf("mockChangeStream: Next past end of mocked iterations (%d iterations)", len(cs.iterations)) } + var remaining time.Duration + if deadline, ok := ctx.Deadline(); ok { + remaining = time.Until(deadline) + } + cs.deadlines = append(cs.deadlines, remaining) + ts := time.Now() cs.emittedTimes = append(cs.emittedTimes, ts) cs.resumeToken = toResumeToken(ts) @@ -66,7 +96,16 @@ func (cs *mockChangeStream) Next(context.Context) bool { switch label { case insert: - cs.current = newInsertChangeEvent(bson.NewObjectID(), ts) + cs.current = newChangeEvent("insert", mockEventID(cs.inserts), ts) + cs.inserts++ + cs.err = nil + return true + case nullIdInsert: + cs.current = newChangeEvent("insert", nil, ts) + cs.err = nil + return true + case unsupportedOp: + cs.current = newChangeEvent("drop", "ev-dropped", ts) cs.err = nil return true case rawEvent: @@ -80,12 +119,60 @@ func (cs *mockChangeStream) Next(context.Context) bool { case idle: cs.err = context.DeadlineExceeded return false + case fatal: + cs.err = errFatalChangeStream + return false default: cs.t.Fatalf("mockChangeStream: unknown label %d", label) return false } } +var errFatalChangeStream = errors.New("mock change stream failure") + +func mockEventID(n int) string { return fmt.Sprintf("ev-%04d", n) } + +// tokenAt is the checkpoint text corresponding to the resume token the mock exposed +// during its idx'th Next() call. +func (cs *mockChangeStream) tokenAt(idx int) string { + cs.t.Helper() + require.Less(cs.t, idx, len(cs.emittedTimes), "change stream never reached iteration %d", idx) + return b64(toResumeToken(cs.emittedTimes[idx])) +} + +// iterationOfToken maps a checkpoint back to the change stream iteration that produced +// it, which is what lets tests compare a committed offset against delivered records. +func (cs *mockChangeStream) iterationOfToken(text string) int { + cs.t.Helper() + for idx := range cs.emittedTimes { + if b64(toResumeToken(cs.emittedTimes[idx])) == text { + return idx + } + } + cs.t.Fatalf("checkpoint %q does not match any resume token the change stream emitted", text) + return -1 +} + +// replicableThrough is the _id of every event emitted at or before iteration idx that +// PullRecords is expected to replicate. An event that can never be decoded is included +// under a sentinel id, so that a checkpoint advancing past it shows up as a violation +// rather than passing silently. +func (cs *mockChangeStream) replicableThrough(idx int) []string { + ids := make([]string, 0, idx+1) + inserts := 0 + for _, iteration := range cs.iterations[:min(idx+1, len(cs.iterations))] { + switch iteration { + case insert: + ids = append(ids, mockEventID(inserts)) + inserts++ + case nullIdInsert: + ids = append(ids, "ev-undecodable") + case idle, unsupportedOp, fatal: + } + } + return ids +} + 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 } @@ -116,13 +203,19 @@ func drainMongoCDCRecordsAsync(t *testing.T, stream *model.CDCStream[model.Recor } func newInsertChangeEvent(id bson.ObjectID, ts time.Time) bson.Raw { + return newChangeEvent("insert", id, ts) +} + +// newChangeEvent builds a change event shaped the way the connector's aggregation +// pipeline projects it. +func newChangeEvent(operationType string, id any, ts time.Time) bson.Raw { wallTime := ts.UTC().Truncate(time.Millisecond) event, _ := bson.Marshal(bson.D{ {Key: "ns", Value: bson.D{ {Key: "db", Value: "db"}, {Key: "coll", Value: "coll"}, }}, - {Key: "operationType", Value: "insert"}, + {Key: "operationType", Value: operationType}, {Key: "documentKey", Value: bson.D{{Key: "_id", Value: id}}}, {Key: "fullDocument", Value: bson.D{ {Key: "_id", Value: id}, diff --git a/flow/go.mod b/flow/go.mod index e0453b7d53..066d0fcdb6 100644 --- a/flow/go.mod +++ b/flow/go.mod @@ -57,6 +57,7 @@ require ( github.com/pingcap/errors v0.11.5-0.20260523003111-3697ad564b43 github.com/pingcap/tidb v0.0.0-20250130070702-43f2fb91d740 github.com/pingcap/tidb/pkg/parser v0.0.0-20260504140133-511dba1dbe17 + github.com/pressly/goose/v3 v3.27.3 github.com/quasilyte/go-ruleguard/dsl v0.3.23 github.com/shopspring/decimal v1.4.0 github.com/slack-go/slack v0.29.0 @@ -85,6 +86,7 @@ require ( go.temporal.io/sdk v1.48.0 go.temporal.io/sdk/contrib/opentelemetry v0.8.1 go.uber.org/automaxprocs v1.6.0 + go.uber.org/goleak v1.3.0 golang.org/x/crypto v0.55.0 golang.org/x/exp v0.0.0-20260824195058-e88cd73687aa golang.org/x/sync v0.22.0 @@ -99,8 +101,6 @@ require ( k8s.io/client-go v0.35.3 // Note: v0.* are newer than v1.* ) -require github.com/pressly/goose/v3 v3.27.3 - require ( cel.dev/expr v0.25.2 // indirect cloud.google.com/go/auth/oauth2adapt v0.2.8 // indirect diff --git a/flow/internal/dynamicconf.go b/flow/internal/dynamicconf.go index 79ce3ab889..ff8c44f4db 100644 --- a/flow/internal/dynamicconf.go +++ b/flow/internal/dynamicconf.go @@ -584,6 +584,14 @@ var DynamicSettings = [...]*protos.DynamicSetting{ ApplyMode: protos.DynconfApplyMode_APPLY_MODE_IMMEDIATE, TargetForSetting: protos.DynconfTarget_POSTGRES, }, + { + Name: "PEERDB_MONGODB_NUM_PARALLEL_DECODE_THREADS", + Description: "Number of parallel threads to use when decoding full BSON documents in MongoDB change events.", + DefaultValue: "1", + ValueType: protos.DynconfValueType_INT, + ApplyMode: protos.DynconfApplyMode_APPLY_MODE_IMMEDIATE, + TargetForSetting: protos.DynconfTarget_ALL, + }, } var DynamicIndex = func() map[string]int { @@ -1032,3 +1040,7 @@ func PeerDBMongoDBExcludedOperationTypes(ctx context.Context, env map[string]str func PeerDBPostgresRawBatchCleanupThreshold(ctx context.Context, env map[string]string) (int64, error) { return dynamicConfSigned[int64](ctx, env, "PEERDB_POSTGRES_RAW_BATCH_CLEANUP_THRESHOLD") } + +func PeerDBMongoDBNumParallelDecodeThreads(ctx context.Context, env map[string]string) (int64, error) { + return dynamicConfSigned[int64](ctx, env, "PEERDB_MONGODB_NUM_PARALLEL_DECODE_THREADS") +} diff --git a/flow/shared/concurrency/workers.go b/flow/shared/concurrency/workers.go new file mode 100644 index 0000000000..d7d60c3222 --- /dev/null +++ b/flow/shared/concurrency/workers.go @@ -0,0 +1,169 @@ +package concurrency + +import ( + "context" + + "golang.org/x/sync/errgroup" +) + +type sendMsg[D, RT any] struct { + decodedChunk D + resumeToken RT +} + +// PullRecordsWorkerPool implements an ordered, batched concurent worker pool for use with +// PullRecords. The canonical use-case is for decoding events (type E) coming from the +// source database into decoded events (type D), with an associated resumeToken +// to thread through to the Send function along with events (type RT). We use the term +// "chunk" to refer to batches of events in this worker, to avoid nomenclature clashes +// elsewhere with batch sizes. +// +// Items of work ("events") arrive through AddItem, and are chunked into instances of WorkerFunc +// that are spun up (up to a max of `Concurrency` in parallel). The result of WorkerFunc +// is passed through to Send, along with the resumeToken from the last call to AddItem +// that added an event that was part of this chunk. Order is maintained through all calls; +// events that arrive earlier will be part of earlier chunks than events that arrived later, +// and chunks will maintain this AddItem-implied order on their way to Send(). Within a chunk, +// the slice of events is also ordered from earliest to latest. +type PullRecordsWorkerPool[E, D, RT any] struct { + WorkerFunc func(events []E) (D, error) + Send func(ctx context.Context, items D, resumeToken RT) error + + lastToken RT + sem chan struct{} + sender chan chan sendMsg[D, RT] + workerCtx context.Context //nolint:containedctx // errgroup ctx must reach the worker in Flush. + ctxCancel context.CancelFunc + eg *errgroup.Group + inProgressChunk []E + + Concurrency, ChunkSize int + + closed bool +} + +// Buffered channel size for channels to pass records to the send function (see below). +// Used as a min; if the passed-in `Concurrency` is higher, that's used instead. +const workerBufferedChanSize = 10 + +// Init initializes this PullRecordsWorker. Once Wait() has returned, Init() +// can be called again to reuse this struct. +func (p *PullRecordsWorkerPool[E, D, RT]) Init(ctx context.Context) { + p.inProgressChunk = make([]E, 0, p.ChunkSize) + p.closed = false + p.sem = make(chan struct{}, max(1, p.Concurrency)) + p.sender = make(chan chan sendMsg[D, RT], max(workerBufferedChanSize, p.Concurrency)) + ctx, p.ctxCancel = context.WithCancel(ctx) //nolint:gosec // G118: cancelled in Wait. + p.eg, ctx = errgroup.WithContext(ctx) + p.workerCtx = ctx + + // Start the send loop. + p.eg.Go(func() error { + for { + select { + case sendChan, ok := <-p.sender: + if !ok { + return nil + } + select { + case msg := <-sendChan: + if err := p.Send(ctx, msg.decodedChunk, msg.resumeToken); err != nil { + return err + } + case <-ctx.Done(): + return ctx.Err() + } + case <-ctx.Done(): + return ctx.Err() + } + } + }) +} + +// AddItem adds one event to be chunked and passed to a worker. Not safe for concurrent +// use. +func (p *PullRecordsWorkerPool[E, D, RT]) AddItem(ctx context.Context, event E, token RT) error { + p.inProgressChunk = append(p.inProgressChunk, event) + p.lastToken = token + if len(p.inProgressChunk) >= p.ChunkSize { + return p.Flush(ctx) + } + return nil +} + +// Flush schedules work for any in-flight events that were added by previous calls to AddItem, +// but were not dispatched to a worker yet. Not safe for concurrent use with AddItem. +func (p *PullRecordsWorkerPool[E, D, RT]) Flush(ctx context.Context) error { + if len(p.inProgressChunk) == 0 { + // Nothing to do. + return nil + } + // Grab a slot in the semaphore. + select { + case p.sem <- struct{}{}: + // Grabbed a slot. + curChunk := p.inProgressChunk + sendChan := make(chan sendMsg[D, RT]) + lastToken := p.lastToken + select { + case p.sender <- sendChan: + case <-ctx.Done(): + p.ctxCancel() + return ctx.Err() + case <-p.workerCtx.Done(): + return nil + } + p.eg.Go(func() error { + defer func() { + // Release the slot. + <-p.sem + }() + + decoded, err := p.WorkerFunc(curChunk) + if err != nil { + return err + } + sendMsg := sendMsg[D, RT]{ + decodedChunk: decoded, + resumeToken: lastToken, + } + select { + case sendChan <- sendMsg: + case <-p.workerCtx.Done(): + return context.Cause(p.workerCtx) + } + return nil + }) + p.inProgressChunk = make([]E, 0, p.ChunkSize) + return nil + case <-ctx.Done(): + p.ctxCancel() + return ctx.Err() + case <-p.workerCtx.Done(): + return context.Cause(p.workerCtx) + } +} + +// Wait waits for all in-progress work to finish, or for the passed-in context to be cancelled, +// whichever happens sooner. +func (p *PullRecordsWorkerPool[E, D, RT]) Wait(ctx context.Context) error { + if p.closed { + // Wait called a second time. + return nil + } + // Gracefully drain. + close(p.sender) + p.closed = true + defer p.ctxCancel() + wgWaiter := make(chan error) + go func() { + wgWaiter <- p.eg.Wait() + }() + select { + case err := <-wgWaiter: + return err + case <-ctx.Done(): + p.ctxCancel() + return <-wgWaiter + } +} diff --git a/flow/shared/concurrency/workers_test.go b/flow/shared/concurrency/workers_test.go new file mode 100644 index 0000000000..64ee7fdb9e --- /dev/null +++ b/flow/shared/concurrency/workers_test.go @@ -0,0 +1,84 @@ +package concurrency + +import ( + "context" + "errors" + "fmt" + "slices" + "testing" + + "github.com/stretchr/testify/require" + "go.uber.org/goleak" +) + +func TestPullRecordsWorkerPoolSendsEveryItemInOrder(t *testing.T) { + var gotChunks [][]int + var gotTokens []string + + pool := PullRecordsWorkerPool[int, []int, string]{ + Concurrency: 4, + ChunkSize: 3, + WorkerFunc: func(events []int) ([]int, error) { return events, nil }, + Send: func(_ context.Context, items []int, resumeToken string) error { + // Send only ever runs on the pool's single send loop, so this needs no locking. + gotChunks = append(gotChunks, items) + gotTokens = append(gotTokens, resumeToken) + return nil + }, + } + pool.Init(t.Context()) + + for i := range 10 { + require.NoError(t, pool.AddItem(t.Context(), i, fmt.Sprintf("token-%d", i))) + } + require.NoError(t, pool.Flush(t.Context())) + require.NoError(t, pool.Wait(t.Context())) + + // 10 items at ChunkSize 3 chunk up as [0 1 2] [3 4 5] [6 7 8] [9], and even with four + // workers decoding in parallel they must reach Send in AddItem order. + require.Equal(t, [][]int{{0, 1, 2}, {3, 4, 5}, {6, 7, 8}, {9}}, gotChunks) + // Each chunk carries the resume token of the last item that went into it. + require.Equal(t, []string{"token-2", "token-5", "token-8", "token-9"}, gotTokens) +} + +func TestPullRecordsWorkerPoolErrorFromWorker(t *testing.T) { + alreadyRunning := goleak.IgnoreCurrent() + defer goleak.VerifyNone(t, alreadyRunning) + + errDecode := errors.New("decode failed") + var lastResumeToken string + barrier := make(chan struct{}, 1) + pool := PullRecordsWorkerPool[int, []int, string]{ + Concurrency: 4, + ChunkSize: 2, + WorkerFunc: func(events []int) ([]int, error) { + if slices.Contains(events, 5) { + // Before returning an error, ensure the previous batch has been sent + // to avoid a flaky test. + <-barrier + return nil, errDecode + } + return events, nil + }, + Send: func(_ context.Context, vals []int, token string) error { + if slices.Contains(vals, 3) { + barrier <- struct{}{} + } + lastResumeToken = token + return nil + }, + } + pool.Init(t.Context()) + + errSeen := false + for i := range 20 { + if err := pool.AddItem(t.Context(), i, fmt.Sprintf("token-%d", i)); err != nil { + require.ErrorIs(t, err, errDecode) + errSeen = true + break + } + } + require.True(t, errSeen, "no error seen from AddItem") + require.Error(t, pool.Wait(t.Context())) + require.Equal(t, "token-3", lastResumeToken) +}