From 08ac10fc1c57559145f33d3babe8f7d13e661b10 Mon Sep 17 00:00:00 2001 From: Bilal Akhtar Date: Fri, 21 Aug 2026 15:53:00 -0400 Subject: [PATCH 1/8] fix(mongo): pipeline document decode in CDC PullRecords Previously, we serially decoded the change event, then the full document that was changed, inside the PullRecords loop for a given batch. This would leave a lot of CPU underutilized as the bulk of work would happen on just one thread. This change moves the full document decode part to a different goroutine, and the record publish to another goroutine, which should allow for greater CPU utilization by the decoding function. --- flow/connectors/mongo/cdc.go | 477 ++++++++++++++++++++---- flow/connectors/mongo/cdc_batch_test.go | 373 ++++++++++++++++++ flow/connectors/mongo/cdc_test.go | 99 ++++- 3 files changed, 864 insertions(+), 85 deletions(-) create mode 100644 flow/connectors/mongo/cdc_batch_test.go diff --git a/flow/connectors/mongo/cdc.go b/flow/connectors/mongo/cdc.go index 203ab911c..b7baa7ca8 100644 --- a/flow/connectors/mongo/cdc.go +++ b/flow/connectors/mongo/cdc.go @@ -6,8 +6,10 @@ import ( "errors" "fmt" "log/slog" + "runtime" "slices" "strings" + "sync" "sync/atomic" "time" @@ -16,6 +18,7 @@ import ( "go.mongodb.org/mongo-driver/v2/mongo/options" "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/trace" + "golang.org/x/sync/errgroup" "github.com/PeerDB-io/peerdb/flow/alerting" "github.com/PeerDB-io/peerdb/flow/generated/protos" @@ -230,12 +233,281 @@ func decodeEvent( return nil } +// Constants used by PullRecords. +// +// Buffered channel size for channels to pass records to decode and send loops (see below). +// This should ideally be larger than decodeWorkerBufSize. +const workerBufferedChanSize = 10 + +// Number of recordItems to pass in one batch to decode/send loops. Doing per-item channel sends +// results in too much coordination and reduces effective concurrency in practice. +const pullRecordsItemsBatchSize = 256 + +// maxNumDecodeWorkers is the maximum number of goroutines that decodeLoop spins up to parallelize +// decoding of record batches. +const maxNumDecodeWorkers = 6 + +// decodeWorkerBufSize is the size of the channel to each decode worker spun up by decodeLoop. +const decodeWorkerBufSize = 4 + +// Amount of time to wait for workers to drain before timing out. +const workerDrainTimeout = 2 * time.Minute + +type decodeBatch struct { + // base64-encoded resume token + resumeToken string + items []recordItems +} + +type sendBatch struct { + // base64-encoded resume token + resumeToken string + // records to send to RecordStream. + records []model.Record[model.RecordItems] +} + +// Two additional loops are created by PullRecords, each running in their own goroutines. +// One is decodeLoop, which takes records from the main PullRecords goroutine through `records` +// that have already had the operation type decoded, but not the full document. decodeLoop manages +// decodeWorkers running in their own goroutines that handle the actual bson document parsing. After +// a decodeWorker has parsed a batch of items, it passes the batch to sendLoop, which assembles records +// from all decodeWorkers into req.RecordStream. req.RecordStream.AddRecord could block on the downstream +// channel send, so separating `sendLoop` out this way from `decodeWorker` makes sense to speed up the +// relatively CPU-heavy `decodeWorker`. `sendLoop` is also responsible for advancing the resume token +// inside req.RecordStream, except in cases when `sendLoop` has been confirmed to be terminated, +// in which case `PullRecords` itself can do the advancement. +// +// decodeLoop is expected to propagate channel closures through to sendLoop for graceful +// draining to ensure no records are lost. Context cancellation is only used in the error +// path, where losing inflight records is acceptable with the expectation that a future +// restart from the last checkpoint will see those events again. Note that in case of +// recreation of the upstream changestream, we could drain and restart these two worker loops +// plus any decodeWorkers within the execution of one `PullRecords`. +func (c *MongoConnector) sendLoop( + ctx context.Context, + records <-chan chan sendBatch, + errChan chan<- error, + req *model.PullRecordsRequest[model.RecordItems], + wg *sync.WaitGroup, +) { + defer wg.Done() + + for { + select { + case recordChan, ok := <-records: + if !ok { + return + } + select { + case sendBatch := <-recordChan: + for i := range sendBatch.records { + if err := req.RecordStream.AddRecord(ctx, sendBatch.records[i]); err != nil { + select { + case errChan <- err: + case <-ctx.Done(): + } + return + } + } + if sendBatch.resumeToken != "" { + req.RecordStream.UpdateLatestCheckpointText(sendBatch.resumeToken) + } + case <-ctx.Done(): + return + } + case <-ctx.Done(): + return + } + } +} + +type recordItems struct { + maybeFullDocument *bson.Raw + operationType operationType + sourceTableName string + destinationTableName string + documentKey bson.Raw + commitTimeNanos int64 +} + +// decodeWorker is spun up by decodeLoop in separate goroutines, up to +// maxNumDecodeWorker in parallel. Batches of items to decode are passed in +// through recv, and the output is sent through `send` to `sendLoop` directly. +func (c *MongoConnector) decodeWorker( + ctx context.Context, + recv <-chan decodeBatch, + send chan<- sendBatch, + req *model.PullRecordsRequest[model.RecordItems], +) error { + // Utils used by this routine. + converter := NewDirectBsonConverter() + fullDocumentColumnName := DefaultFullDocumentColumnName + if req.InternalVersion < shared.InternalVersion_MongoDBFullDocumentColumnToDoc { + fullDocumentColumnName = LegacyFullDocumentColumnName + } + parseItem := func(item recordItems) (model.Record[model.RecordItems], error) { + items := model.NewRecordItems(2) + + if len(item.documentKey) > 0 { + rv := item.documentKey.Lookup(DefaultDocumentKeyColumnName) + if rv.IsZero() || rv.Type == bson.TypeNull { + return nil, exceptions.NewInvalidIdValueError(item.sourceTableName) + } + qValue, err := converter.QValueStringFromId(rv, req.InternalVersion) + if err != nil { + return nil, err + } + items.AddColumn(DefaultDocumentKeyColumnName, qValue) + } else { + return nil, fmt.Errorf("document key is nil") + } + + if item.maybeFullDocument != nil && len(*item.maybeFullDocument) > 0 { + qValue, err := converter.QValueJSONFromDocument(*item.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 item.operationType { + case operationTypeInsert: + record = &model.InsertRecord[model.RecordItems]{ + BaseRecord: model.BaseRecord{CommitTimeNano: item.commitTimeNanos}, + Items: items, + SourceTableName: item.sourceTableName, + DestinationTableName: item.destinationTableName, + } + + case operationTypeUpdate, operationTypeReplace: + record = &model.UpdateRecord[model.RecordItems]{ + BaseRecord: model.BaseRecord{CommitTimeNano: item.commitTimeNanos}, + NewItems: items, + SourceTableName: item.sourceTableName, + DestinationTableName: item.destinationTableName, + } + case operationTypeDelete: + record = &model.DeleteRecord[model.RecordItems]{ + BaseRecord: model.BaseRecord{CommitTimeNano: item.commitTimeNanos}, + Items: items, + SourceTableName: item.sourceTableName, + DestinationTableName: item.destinationTableName, + } + } + return record, nil + } + for { + select { + case batch, ok := <-recv: + if !ok { + return nil + } + items := batch.items + + modelRecords := make([]model.Record[model.RecordItems], len(items)) + for i := range items { + var err error + if modelRecords[i], err = parseItem(items[i]); err != nil { + return err + } + } + + select { + case send <- sendBatch{records: modelRecords, resumeToken: batch.resumeToken}: + case <-ctx.Done(): + return ctx.Err() + } + case <-ctx.Done(): + return ctx.Err() + } + } +} + +func (c *MongoConnector) decodeLoop( + parentCtx context.Context, + recv <-chan decodeBatch, + errChan chan<- error, + req *model.PullRecordsRequest[model.RecordItems], + wg *sync.WaitGroup, +) { + defer wg.Done() + + // NB: child workers are tracked by this errGroup, not by + // the parent WaitGroup (shared with PullRecords and sendLoop). + eg, ctx := errgroup.WithContext(parentCtx) + defer func() { + err := eg.Wait() + if err != nil { + select { + case errChan <- err: + case <-parentCtx.Done(): + } + } + }() + + // This function spins up half as many workers as the number of cores, capped + // at maxNumDecodeWorkers. + numWorkers := max(1, min(maxNumDecodeWorkers, runtime.GOMAXPROCS(0)/2)) + workerChan := make([]struct { + req chan decodeBatch + res chan sendBatch + }, numWorkers) + for i := range workerChan { + workerChan[i].req = make(chan decodeBatch, decodeWorkerBufSize) + workerChan[i].res = make(chan sendBatch, decodeWorkerBufSize) + // Start worker. + eg.Go(func() error { + return c.decodeWorker(ctx, workerChan[i].req, workerChan[i].res, req) + }) + } + + // Start up sendLoop. Have a buffered channel in case decoding runs faster + // than the downstream addition of records to req.RecordStream. + sender := make(chan chan sendBatch, decodeWorkerBufSize*maxNumDecodeWorkers) + go c.sendLoop(parentCtx, sender, errChan, req, wg) + // Index of next worker to send a result to. + nextWorker := 0 + + for { + select { + case item, ok := <-recv: + if !ok { + // Draining. Close all worker channels. + for i := range workerChan { + close(workerChan[i].req) + } + close(sender) + return + } + + select { + case workerChan[nextWorker].req <- item: + sender <- workerChan[nextWorker].res + nextWorker = (nextWorker + 1) % numWorkers + case <-ctx.Done(): + return + } + case <-parentCtx.Done(): + // Context cancellation usually means an error, not just cutting a batch (which would be + // a graceful drain via recv getting closed). + return + } + } +} + func (c *MongoConnector) PullRecords( ctx context.Context, catalogPool shared.CatalogPool, otelManager *otel_metrics.OtelManager, req *model.PullRecordsRequest[model.RecordItems], -) error { +) (retErr error) { defer req.RecordStream.Close() var alerter *alerting.Alerter @@ -247,6 +519,7 @@ func (c *MongoConnector) PullRecords( if req.InternalVersion < shared.InternalVersion_MongoDBFullDocumentColumnToDoc { fullDocumentColumnName = LegacyFullDocumentColumnName } + var wg sync.WaitGroup c.logger.Info("[mongo] started PullRecords for mirror "+req.FlowJobName, slog.Any("table_mapping", req.TableNameMapping), @@ -324,7 +597,12 @@ func (c *MongoConnector) PullRecords( attribute.Int64(otel_metrics.RowsInBatchKey, int64(recordCount)), attribute.Int64(otel_metrics.BytesPulledKey, cumulativeBytesProcessed.Load()), ) - if changeStream != nil { + if changeStream != nil && retErr == nil { + // NB: We don't set the otel metrics ResumeTokenKey in the error case because + // it's possible for the changeStream to have been advanced past the last record + // send to RecordStream (and therefore the actual resume token set in RecordStream). + // It's safer to not report a resume token than to report an incorrect one that can + // skip records if used. if rt := changeStream.ResumeToken(); rt != nil { rtStr := base64.StdEncoding.EncodeToString(rt) if len(rtStr) > 64 { @@ -339,6 +617,15 @@ func (c *MongoConnector) PullRecords( slog.Int("channelLen", req.RecordStream.ChannelLen()), slog.Float64("elapsedMinutes", time.Since(pullStart).Minutes())) }() + // Context inheritance tree: There's a parent ctx (always called ctx in this function). + // Off of that, we create two child contexts, one for workers (called workerCtx), that's passed + // to children workers (eg. decodeLoop). The other, timeoutCtx, is managed by us for timeouts. + // Note that when we hit a timeout, we don't want to cancel the workerCtx; we want the workers + // to gracefully drain. + // + // Also note that timeoutCtx is occasionally rewritten, such as when a record in a batch arrives + // or when we reset the changestream. + workerCtx, workerCtxCancel := context.WithCancel(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 +638,8 @@ func (c *MongoConnector) PullRecords( defer func() { cancelTimeout() + workerCtxCancel() + wg.Wait() reportBytesShutdown() read := deltaBytesProcessed.Swap(0) otelManager.Metrics.FetchedBytesCounter.Add(ctx, read) @@ -377,44 +666,8 @@ 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 @@ -426,9 +679,69 @@ func (c *MongoConnector) PullRecords( slog.Int("channelLen", req.RecordStream.ChannelLen()), slog.Float64("elapsedMinutes", time.Since(pullStart).Minutes())) } + } + existingBatch := make([]recordItems, 0, pullRecordsItemsBatchSize) + decodeChan := make(chan decodeBatch, workerBufferedChanSize) + errChan := make(chan error, 2) // size = one for decodeLoop, one for sendLoop. + // decodeLoop starts up sendLoop. + wg.Add(2) + go c.decodeLoop(workerCtx, decodeChan, errChan, req, &wg) + + finishBatch := func() error { + if len(existingBatch) == 0 { + // Nothing to send. + return nil + } + 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) + } + select { + case decodeChan <- decodeBatch{items: existingBatch, resumeToken: rtText}: + case err := <-errChan: + workerCtxCancel() + return err + case <-ctx.Done(): + workerCtxCancel() + return ctx.Err() + } + existingBatch = make([]recordItems, 0, pullRecordsItemsBatchSize) return nil } + drainWorkers := func() error { + finishErr := finishBatch() + // Close the decode loop and wait for events to drain. + close(decodeChan) + wgWaiter := make(chan bool) + go func() { + wg.Wait() + close(wgWaiter) + }() + select { + case <-wgWaiter: + // Check if there was an error returned at the same time + // as the other goroutines disappeared. + select { + case err := <-errChan: + return err + default: + } + case err := <-errChan: + workerCtxCancel() + <-wgWaiter + return err + case <-time.After(workerDrainTimeout): + workerCtxCancel() + <-wgWaiter + return errors.New("timed out waiting for PullRecords workers to drain") + } + return finishErr + } + recreateChangeStream := func(useOperationTime bool) error { // extract the most recent resumeToken resumeToken := changeStream.ResumeToken() @@ -445,6 +758,12 @@ func (c *MongoConnector) PullRecords( cancelTimeout() timeoutCtx, cancelTimeout = context.WithTimeout(ctx, time.Hour) + // Restart the decode loop + decodeChan = make(chan decodeBatch, workerBufferedChanSize) + // decodeLoop starts up sendLoop; add one for each. + wg.Add(2) + go c.decodeLoop(workerCtx, decodeChan, errChan, req, &wg) + // set resume point based on whether operation time should be used or not if useOperationTime { timestamp, err := decodeTimestampFromResumeToken(resumeToken) @@ -474,11 +793,25 @@ func (c *MongoConnector) PullRecords( return fmt.Errorf("unexpected: changestream.Next() returned false but no change stream error was recorded") } + // Before checking for timeout, see if there's an error from the child workers waiting. + select { + case err := <-errChan: + workerCtxCancel() + return err + default: + } + if err := drainWorkers(); 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 +872,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) - } + items := recordItems{ + documentKey: changeEvent.DocumentKey, + maybeFullDocument: changeEvent.FullDocument, + operationType: operationType(changeEvent.OperationType), + sourceTableName: sourceTableName, + destinationTableName: destinationTableName, + commitTimeNanos: commitTimeNanos, + } + switch items.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 +904,16 @@ func (c *MongoConnector) PullRecords( continue } otelManager.Metrics.FetchedEventSizeHistogram.Record(ctx, changeEventSize) - checkpoint() + existingBatch = append(existingBatch, items) + if len(existingBatch) >= pullRecordsItemsBatchSize { + if err := finishBatch(); err != nil { + return err + } + } + } + + if err := drainWorkers(); 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 000000000..b5086e3be --- /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(pullRecordsItemsBatchSize + 9), + maxBatchSize: pullRecordsItemsBatchSize + 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(pullRecordsItemsBatchSize), nullIdInsert), + maxBatchSize: pullRecordsItemsBatchSize + 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 pullRecordsItemsBatchSize 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*pullRecordsItemsBatchSize + 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*pullRecordsItemsBatchSize)...) + h := newPullHarness(t, iterations...) + h.req.MaxBatchSize = uint32(2 * pullRecordsItemsBatchSize) + + 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*pullRecordsItemsBatchSize)...) + h.req.MaxBatchSize = 3 * pullRecordsItemsBatchSize + // 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*pullRecordsItemsBatchSize { + 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 51e7a0cfa..797f75b26 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}, From f9108b3a9271aecc3752d2de38ca0134dd890b24 Mon Sep 17 00:00:00 2001 From: Bilal Akhtar Date: Wed, 2 Sep 2026 16:32:54 -0400 Subject: [PATCH 2/8] implement CR suggestions, simplify concurrency handling --- flow/connectors/mongo/cdc.go | 299 ++++++++++------------------------- 1 file changed, 87 insertions(+), 212 deletions(-) diff --git a/flow/connectors/mongo/cdc.go b/flow/connectors/mongo/cdc.go index b7baa7ca8..4dfaad524 100644 --- a/flow/connectors/mongo/cdc.go +++ b/flow/connectors/mongo/cdc.go @@ -9,7 +9,6 @@ import ( "runtime" "slices" "strings" - "sync" "sync/atomic" "time" @@ -239,84 +238,59 @@ func decodeEvent( // This should ideally be larger than decodeWorkerBufSize. const workerBufferedChanSize = 10 -// Number of recordItems to pass in one batch to decode/send loops. Doing per-item channel sends +// Number of recordItems to pass in one chunk to decode/send loops. Doing per-item channel sends // results in too much coordination and reduces effective concurrency in practice. -const pullRecordsItemsBatchSize = 256 +const pullRecordsItemsChunkSize = 256 // maxNumDecodeWorkers is the maximum number of goroutines that decodeLoop spins up to parallelize -// decoding of record batches. +// decoding of record chunks. const maxNumDecodeWorkers = 6 -// decodeWorkerBufSize is the size of the channel to each decode worker spun up by decodeLoop. -const decodeWorkerBufSize = 4 - -// Amount of time to wait for workers to drain before timing out. -const workerDrainTimeout = 2 * time.Minute - -type decodeBatch struct { +type decodeChunk struct { // base64-encoded resume token resumeToken string items []recordItems } -type sendBatch struct { +type sendChunk struct { // base64-encoded resume token resumeToken string // records to send to RecordStream. records []model.Record[model.RecordItems] } -// Two additional loops are created by PullRecords, each running in their own goroutines. -// One is decodeLoop, which takes records from the main PullRecords goroutine through `records` -// that have already had the operation type decoded, but not the full document. decodeLoop manages -// decodeWorkers running in their own goroutines that handle the actual bson document parsing. After -// a decodeWorker has parsed a batch of items, it passes the batch to sendLoop, which assembles records -// from all decodeWorkers into req.RecordStream. req.RecordStream.AddRecord could block on the downstream -// channel send, so separating `sendLoop` out this way from `decodeWorker` makes sense to speed up the -// relatively CPU-heavy `decodeWorker`. `sendLoop` is also responsible for advancing the resume token -// inside req.RecordStream, except in cases when `sendLoop` has been confirmed to be terminated, -// in which case `PullRecords` itself can do the advancement. -// -// decodeLoop is expected to propagate channel closures through to sendLoop for graceful -// draining to ensure no records are lost. Context cancellation is only used in the error -// path, where losing inflight records is acceptable with the expectation that a future -// restart from the last checkpoint will see those events again. Note that in case of -// recreation of the upstream changestream, we could drain and restart these two worker loops -// plus any decodeWorkers within the execution of one `PullRecords`. +// PullRecords spins up two kinds of goroutines: one sendLoop that lives for the duration of one +// ChangeStream and is responsible for updating the latest checkpoint info. The other type is implemented +// in decodeWorker and is spun up spun up per-chunk (with parallelism managed by decodeWorkerSem), to +// decode one chunk at a time. The context passed into these goroutines is shared, but is not used +// to signal timeouts; rather, when cutting a batch, the `records` chan passed to sendLoop below is +// closed, and we wait for all child goroutines to gracefully drain. func (c *MongoConnector) sendLoop( ctx context.Context, - records <-chan chan sendBatch, - errChan chan<- error, + records <-chan chan sendChunk, req *model.PullRecordsRequest[model.RecordItems], - wg *sync.WaitGroup, -) { - defer wg.Done() - +) error { for { select { case recordChan, ok := <-records: if !ok { - return + return nil } select { - case sendBatch := <-recordChan: - for i := range sendBatch.records { - if err := req.RecordStream.AddRecord(ctx, sendBatch.records[i]); err != nil { - select { - case errChan <- err: - case <-ctx.Done(): - } - return + case sendChunk := <-recordChan: + for i := range sendChunk.records { + if err := req.RecordStream.AddRecord(ctx, sendChunk.records[i]); err != nil { + return err } } - if sendBatch.resumeToken != "" { - req.RecordStream.UpdateLatestCheckpointText(sendBatch.resumeToken) + if sendChunk.resumeToken != "" { + req.RecordStream.UpdateLatestCheckpointText(sendChunk.resumeToken) } case <-ctx.Done(): - return + return nil } case <-ctx.Done(): - return + return nil } } } @@ -330,13 +304,13 @@ type recordItems struct { commitTimeNanos int64 } -// decodeWorker is spun up by decodeLoop in separate goroutines, up to -// maxNumDecodeWorker in parallel. Batches of items to decode are passed in -// through recv, and the output is sent through `send` to `sendLoop` directly. +// decodeWorker is spun up by PullRecords in separate goroutines, up to +// maxNumDecodeWorker in parallel. The output is sent through `send` to `sendLoop` +// directly. func (c *MongoConnector) decodeWorker( ctx context.Context, - recv <-chan decodeBatch, - send chan<- sendBatch, + chunk decodeChunk, + send chan<- sendChunk, req *model.PullRecordsRequest[model.RecordItems], ) error { // Utils used by this routine. @@ -355,7 +329,7 @@ func (c *MongoConnector) decodeWorker( } qValue, err := converter.QValueStringFromId(rv, req.InternalVersion) if err != nil { - return nil, err + return nil, fmt.Errorf("failed to convert key: %w", err) } items.AddColumn(DefaultDocumentKeyColumnName, qValue) } else { @@ -403,103 +377,19 @@ func (c *MongoConnector) decodeWorker( } return record, nil } - for { - select { - case batch, ok := <-recv: - if !ok { - return nil - } - items := batch.items - - modelRecords := make([]model.Record[model.RecordItems], len(items)) - for i := range items { - var err error - if modelRecords[i], err = parseItem(items[i]); err != nil { - return err - } - } - - select { - case send <- sendBatch{records: modelRecords, resumeToken: batch.resumeToken}: - case <-ctx.Done(): - return ctx.Err() - } - case <-ctx.Done(): - return ctx.Err() - } - } -} - -func (c *MongoConnector) decodeLoop( - parentCtx context.Context, - recv <-chan decodeBatch, - errChan chan<- error, - req *model.PullRecordsRequest[model.RecordItems], - wg *sync.WaitGroup, -) { - defer wg.Done() - - // NB: child workers are tracked by this errGroup, not by - // the parent WaitGroup (shared with PullRecords and sendLoop). - eg, ctx := errgroup.WithContext(parentCtx) - defer func() { - err := eg.Wait() - if err != nil { - select { - case errChan <- err: - case <-parentCtx.Done(): - } + modelRecords := make([]model.Record[model.RecordItems], len(chunk.items)) + for i := range chunk.items { + var err error + if modelRecords[i], err = parseItem(chunk.items[i]); err != nil { + return err } - }() - - // This function spins up half as many workers as the number of cores, capped - // at maxNumDecodeWorkers. - numWorkers := max(1, min(maxNumDecodeWorkers, runtime.GOMAXPROCS(0)/2)) - workerChan := make([]struct { - req chan decodeBatch - res chan sendBatch - }, numWorkers) - for i := range workerChan { - workerChan[i].req = make(chan decodeBatch, decodeWorkerBufSize) - workerChan[i].res = make(chan sendBatch, decodeWorkerBufSize) - // Start worker. - eg.Go(func() error { - return c.decodeWorker(ctx, workerChan[i].req, workerChan[i].res, req) - }) } - - // Start up sendLoop. Have a buffered channel in case decoding runs faster - // than the downstream addition of records to req.RecordStream. - sender := make(chan chan sendBatch, decodeWorkerBufSize*maxNumDecodeWorkers) - go c.sendLoop(parentCtx, sender, errChan, req, wg) - // Index of next worker to send a result to. - nextWorker := 0 - - for { - select { - case item, ok := <-recv: - if !ok { - // Draining. Close all worker channels. - for i := range workerChan { - close(workerChan[i].req) - } - close(sender) - return - } - - select { - case workerChan[nextWorker].req <- item: - sender <- workerChan[nextWorker].res - nextWorker = (nextWorker + 1) % numWorkers - case <-ctx.Done(): - return - } - case <-parentCtx.Done(): - // Context cancellation usually means an error, not just cutting a batch (which would be - // a graceful drain via recv getting closed). - return - } + select { + case send <- sendChunk{records: modelRecords, resumeToken: chunk.resumeToken}: + case <-ctx.Done(): + return ctx.Err() } + return nil } func (c *MongoConnector) PullRecords( @@ -507,7 +397,7 @@ func (c *MongoConnector) PullRecords( catalogPool shared.CatalogPool, otelManager *otel_metrics.OtelManager, req *model.PullRecordsRequest[model.RecordItems], -) (retErr error) { +) error { defer req.RecordStream.Close() var alerter *alerting.Alerter @@ -519,7 +409,6 @@ func (c *MongoConnector) PullRecords( if req.InternalVersion < shared.InternalVersion_MongoDBFullDocumentColumnToDoc { fullDocumentColumnName = LegacyFullDocumentColumnName } - var wg sync.WaitGroup c.logger.Info("[mongo] started PullRecords for mirror "+req.FlowJobName, slog.Any("table_mapping", req.TableNameMapping), @@ -597,12 +486,7 @@ func (c *MongoConnector) PullRecords( attribute.Int64(otel_metrics.RowsInBatchKey, int64(recordCount)), attribute.Int64(otel_metrics.BytesPulledKey, cumulativeBytesProcessed.Load()), ) - if changeStream != nil && retErr == nil { - // NB: We don't set the otel metrics ResumeTokenKey in the error case because - // it's possible for the changeStream to have been advanced past the last record - // send to RecordStream (and therefore the actual resume token set in RecordStream). - // It's safer to not report a resume token than to report an incorrect one that can - // skip records if used. + if changeStream != nil { if rt := changeStream.ResumeToken(); rt != nil { rtStr := base64.StdEncoding.EncodeToString(rt) if len(rtStr) > 64 { @@ -619,16 +503,18 @@ func (c *MongoConnector) PullRecords( }() // Context inheritance tree: There's a parent ctx (always called ctx in this function). // Off of that, we create two child contexts, one for workers (called workerCtx), that's passed - // to children workers (eg. decodeLoop). The other, timeoutCtx, is managed by us for timeouts. + // to children goroutines (eg. decodeWorker). The other, timeoutCtx, is managed by us for timeouts. // Note that when we hit a timeout, we don't want to cancel the workerCtx; we want the workers // to gracefully drain. // - // Also note that timeoutCtx is occasionally rewritten, such as when a record in a batch arrives + // Also note that timeoutCtx is occasionally rewritten, such as when the first record arrives, // or when we reset the changestream. workerCtx, workerCtxCancel := context.WithCancel(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) + var workerEg *errgroup.Group + workerEg, workerCtx = errgroup.WithContext(workerCtx) reportBytesShutdown := common.Interval(ctx, time.Second*10, func() { read := deltaBytesProcessed.Swap(0) @@ -639,7 +525,7 @@ func (c *MongoConnector) PullRecords( defer func() { cancelTimeout() workerCtxCancel() - wg.Wait() + workerEg.Wait() reportBytesShutdown() read := deltaBytesProcessed.Swap(0) otelManager.Metrics.FetchedBytesCounter.Add(ctx, read) @@ -680,15 +566,16 @@ func (c *MongoConnector) PullRecords( slog.Float64("elapsedMinutes", time.Since(pullStart).Minutes())) } } - existingBatch := make([]recordItems, 0, pullRecordsItemsBatchSize) - decodeChan := make(chan decodeBatch, workerBufferedChanSize) - errChan := make(chan error, 2) // size = one for decodeLoop, one for sendLoop. - // decodeLoop starts up sendLoop. - wg.Add(2) - go c.decodeLoop(workerCtx, decodeChan, errChan, req, &wg) - - finishBatch := func() error { - if len(existingBatch) == 0 { + existingChunk := make([]recordItems, 0, pullRecordsItemsChunkSize) + sendChan := make(chan chan sendChunk, workerBufferedChanSize) + numDecodeWorkers := max(1, min(maxNumDecodeWorkers, runtime.GOMAXPROCS(0)/2)) + decodeWorkerSem := make(chan struct{}, numDecodeWorkers) + workerEg.Go(func() error { + return c.sendLoop(workerCtx, sendChan, req) + }) + + dispatchChunk := func() error { + if len(existingChunk) == 0 { // Nothing to send. return nil } @@ -700,44 +587,39 @@ func (c *MongoConnector) PullRecords( rtText = base64.StdEncoding.EncodeToString(rt) } select { - case decodeChan <- decodeBatch{items: existingBatch, resumeToken: rtText}: - case err := <-errChan: - workerCtxCancel() - return err + case decodeWorkerSem <- struct{}{}: case <-ctx.Done(): workerCtxCancel() return ctx.Err() + case <-workerCtx.Done(): + return workerCtx.Err() } - existingBatch = make([]recordItems, 0, pullRecordsItemsBatchSize) + resultChan := make(chan sendChunk) + select { + case sendChan <- resultChan: + case <-ctx.Done(): + workerCtxCancel() + return ctx.Err() + case <-workerCtx.Done(): + return workerCtx.Err() + } + workerEg.Go(func() error { + defer func() { + <-decodeWorkerSem + }() + return c.decodeWorker( + workerCtx, decodeChunk{items: existingChunk, resumeToken: rtText}, resultChan, req) + }) + existingChunk = make([]recordItems, 0, pullRecordsItemsChunkSize) return nil } drainWorkers := func() error { - finishErr := finishBatch() - // Close the decode loop and wait for events to drain. - close(decodeChan) - wgWaiter := make(chan bool) - go func() { - wg.Wait() - close(wgWaiter) - }() - select { - case <-wgWaiter: - // Check if there was an error returned at the same time - // as the other goroutines disappeared. - select { - case err := <-errChan: - return err - default: - } - case err := <-errChan: - workerCtxCancel() - <-wgWaiter + finishErr := dispatchChunk() + // Close the sender chan and wait for events to drain. + close(sendChan) + if err := workerEg.Wait(); err != nil { return err - case <-time.After(workerDrainTimeout): - workerCtxCancel() - <-wgWaiter - return errors.New("timed out waiting for PullRecords workers to drain") } return finishErr } @@ -758,11 +640,11 @@ func (c *MongoConnector) PullRecords( cancelTimeout() timeoutCtx, cancelTimeout = context.WithTimeout(ctx, time.Hour) - // Restart the decode loop - decodeChan = make(chan decodeBatch, workerBufferedChanSize) - // decodeLoop starts up sendLoop; add one for each. - wg.Add(2) - go c.decodeLoop(workerCtx, decodeChan, errChan, req, &wg) + // Restart the send loop. + sendChan = make(chan chan sendChunk, workerBufferedChanSize) + workerEg.Go(func() error { + return c.sendLoop(workerCtx, sendChan, req) + }) // set resume point based on whether operation time should be used or not if useOperationTime { @@ -793,13 +675,6 @@ func (c *MongoConnector) PullRecords( return fmt.Errorf("unexpected: changestream.Next() returned false but no change stream error was recorded") } - // Before checking for timeout, see if there's an error from the child workers waiting. - select { - case err := <-errChan: - workerCtxCancel() - return err - default: - } if err := drainWorkers(); err != nil { return err } @@ -904,9 +779,9 @@ func (c *MongoConnector) PullRecords( continue } otelManager.Metrics.FetchedEventSizeHistogram.Record(ctx, changeEventSize) - existingBatch = append(existingBatch, items) - if len(existingBatch) >= pullRecordsItemsBatchSize { - if err := finishBatch(); err != nil { + existingChunk = append(existingChunk, items) + if len(existingChunk) >= pullRecordsItemsChunkSize { + if err := dispatchChunk(); err != nil { return err } } From 4ceee37e0171757bdc35e4fc65948414beff7216 Mon Sep 17 00:00:00 2001 From: Bilal Akhtar Date: Wed, 2 Sep 2026 21:13:15 -0400 Subject: [PATCH 3/8] add dynamic conf --- flow/connectors/mongo/cdc.go | 16 +++++++--------- flow/connectors/mongo/cdc_batch_test.go | 22 +++++++++++----------- flow/internal/dynamicconf.go | 12 ++++++++++++ 3 files changed, 30 insertions(+), 20 deletions(-) diff --git a/flow/connectors/mongo/cdc.go b/flow/connectors/mongo/cdc.go index 4dfaad524..97ad7878f 100644 --- a/flow/connectors/mongo/cdc.go +++ b/flow/connectors/mongo/cdc.go @@ -6,7 +6,6 @@ import ( "errors" "fmt" "log/slog" - "runtime" "slices" "strings" "sync/atomic" @@ -242,10 +241,6 @@ const workerBufferedChanSize = 10 // results in too much coordination and reduces effective concurrency in practice. const pullRecordsItemsChunkSize = 256 -// maxNumDecodeWorkers is the maximum number of goroutines that decodeLoop spins up to parallelize -// decoding of record chunks. -const maxNumDecodeWorkers = 6 - type decodeChunk struct { // base64-encoded resume token resumeToken string @@ -305,8 +300,8 @@ type recordItems struct { } // decodeWorker is spun up by PullRecords in separate goroutines, up to -// maxNumDecodeWorker in parallel. The output is sent through `send` to `sendLoop` -// directly. +// PEERDB_MONGODB_NUM_PARALLEL_DECODE_THREADS in parallel. The output is sent through `send` to +// `sendLoop` directly. func (c *MongoConnector) decodeWorker( ctx context.Context, chunk decodeChunk, @@ -568,8 +563,11 @@ func (c *MongoConnector) PullRecords( } existingChunk := make([]recordItems, 0, pullRecordsItemsChunkSize) sendChan := make(chan chan sendChunk, workerBufferedChanSize) - numDecodeWorkers := max(1, min(maxNumDecodeWorkers, runtime.GOMAXPROCS(0)/2)) - decodeWorkerSem := make(chan struct{}, numDecodeWorkers) + numParallelDecodeWorkers, err := internal.PeerDBMongoDBNumParallelDecodeThreads(ctx, req.Env) + if err != nil { + return err + } + decodeWorkerSem := make(chan struct{}, max(1, numParallelDecodeWorkers)) workerEg.Go(func() error { return c.sendLoop(workerCtx, sendChan, req) }) diff --git a/flow/connectors/mongo/cdc_batch_test.go b/flow/connectors/mongo/cdc_batch_test.go index b5086e3be..6227c9515 100644 --- a/flow/connectors/mongo/cdc_batch_test.go +++ b/flow/connectors/mongo/cdc_batch_test.go @@ -172,8 +172,8 @@ func TestPullRecordsOffsetNeverRunsAheadOfDeliveredRecords(t *testing.T) { {name: "cut by max batch size", iterations: repeatInserts(8), maxBatchSize: 5}, { name: "cut by max batch size across decode batches", - iterations: repeatInserts(pullRecordsItemsBatchSize + 9), - maxBatchSize: pullRecordsItemsBatchSize + 3, + 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}, @@ -183,8 +183,8 @@ func TestPullRecordsOffsetNeverRunsAheadOfDeliveredRecords(t *testing.T) { {name: "undecodable document", iterations: []iterationType{insert, insert, nullIdInsert}, maxBatchSize: 3}, { name: "undecodable document after a full decode batch", - iterations: append(repeatInserts(pullRecordsItemsBatchSize), nullIdInsert), - maxBatchSize: pullRecordsItemsBatchSize + 1, + 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. @@ -221,11 +221,11 @@ func TestPullRecordsTruncatesBatchAtMaxBatchSize(t *testing.T) { h.requireOffsetsCoverOnlyDeliveredRecords(t, out) } -// A batch larger than pullRecordsItemsBatchSize is handed to the decode workers in +// 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*pullRecordsItemsBatchSize + 37 + maxBatchSize := 2*pullRecordsItemsChunkSize + 37 h := newPullHarness(t, repeatInserts(maxBatchSize+1)...) h.req.MaxBatchSize = uint32(maxBatchSize) @@ -340,9 +340,9 @@ func TestPullRecordsUndecodableDocumentFailsPull(t *testing.T) { 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*pullRecordsItemsBatchSize)...) + iterations := append([]iterationType{nullIdInsert}, repeatInserts(3*pullRecordsItemsChunkSize)...) h := newPullHarness(t, iterations...) - h.req.MaxBatchSize = uint32(2 * pullRecordsItemsBatchSize) + 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") @@ -353,12 +353,12 @@ func TestPullRecordsUndecodableDocumentInEarlierSubBatchFailsPull(t *testing.T) // 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*pullRecordsItemsBatchSize)...) - h.req.MaxBatchSize = 3 * pullRecordsItemsBatchSize + 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*pullRecordsItemsBatchSize { + if idx == 2*pullRecordsItemsChunkSize { h.cancel() } } diff --git a/flow/internal/dynamicconf.go b/flow/internal/dynamicconf.go index 79ce3ab88..ff8c44f4d 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") +} From 080305a796557a916ff219c529cefc2faec4a7c2 Mon Sep 17 00:00:00 2001 From: Bilal Akhtar Date: Thu, 3 Sep 2026 15:17:49 -0400 Subject: [PATCH 4/8] bugfixes in refactor identified by tests --- flow/connectors/mongo/cdc.go | 59 ++++++++++++++++++++---------------- 1 file changed, 33 insertions(+), 26 deletions(-) diff --git a/flow/connectors/mongo/cdc.go b/flow/connectors/mongo/cdc.go index 97ad7878f..1dd63161b 100644 --- a/flow/connectors/mongo/cdc.go +++ b/flow/connectors/mongo/cdc.go @@ -264,6 +264,7 @@ func (c *MongoConnector) sendLoop( ctx context.Context, records <-chan chan sendChunk, req *model.PullRecordsRequest[model.RecordItems], + signalledAsNonEmpty *bool, ) error { for { select { @@ -274,6 +275,13 @@ func (c *MongoConnector) sendLoop( select { case sendChunk := <-recordChan: for i := range sendChunk.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, sendChunk.records[i]); err != nil { return err } @@ -400,11 +408,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)), @@ -471,6 +474,7 @@ func (c *MongoConnector) PullRecords( var recordCount uint32 var deltaBytesProcessed, cumulativeBytesProcessed atomic.Int64 + var signalledAsNonEmpty bool pullStart := time.Now() defer func() { if recordCount == 0 { @@ -500,7 +504,7 @@ func (c *MongoConnector) PullRecords( // Off of that, we create two child contexts, one for workers (called workerCtx), that's passed // to children goroutines (eg. decodeWorker). The other, timeoutCtx, is managed by us for timeouts. // Note that when we hit a timeout, we don't want to cancel the workerCtx; we want the workers - // to gracefully drain. + // to gracefully drain. We recreate workerCtx and workerEg anytime we recreate the changestream. // // Also note that timeoutCtx is occasionally rewritten, such as when the first record arrives, // or when we reset the changestream. @@ -520,7 +524,7 @@ func (c *MongoConnector) PullRecords( defer func() { cancelTimeout() workerCtxCancel() - workerEg.Wait() + _ = workerEg.Wait() reportBytesShutdown() read := deltaBytesProcessed.Swap(0) otelManager.Metrics.FetchedBytesCounter.Add(ctx, read) @@ -550,7 +554,6 @@ func (c *MongoConnector) PullRecords( incrementRecordCount := func() { recordCount += 1 if recordCount == 1 { - req.RecordStream.SignalAsNotEmpty() timeoutCtx, cancelTimeout = context.WithTimeout(ctx, req.IdleTimeout) //nolint:gosec // G118: cancelTimeout called in defer } if recordCount%50000 == 0 { @@ -569,7 +572,7 @@ func (c *MongoConnector) PullRecords( } decodeWorkerSem := make(chan struct{}, max(1, numParallelDecodeWorkers)) workerEg.Go(func() error { - return c.sendLoop(workerCtx, sendChan, req) + return c.sendLoop(workerCtx, sendChan, req, &signalledAsNonEmpty) }) dispatchChunk := func() error { @@ -586,28 +589,26 @@ func (c *MongoConnector) PullRecords( } select { case decodeWorkerSem <- struct{}{}: - case <-ctx.Done(): - workerCtxCancel() - return ctx.Err() case <-workerCtx.Done(): - return workerCtx.Err() + return workerEg.Wait() } resultChan := make(chan sendChunk) select { case sendChan <- resultChan: - case <-ctx.Done(): - workerCtxCancel() - return ctx.Err() case <-workerCtx.Done(): - return workerCtx.Err() - } - workerEg.Go(func() error { - defer func() { - <-decodeWorkerSem - }() - return c.decodeWorker( - workerCtx, decodeChunk{items: existingChunk, resumeToken: rtText}, resultChan, req) - }) + return workerEg.Wait() + } + workerEg.Go(func(existingChunk []recordItems) func() error { + // Two-level method to capture existingChunk before it changes + // below. + return func() error { + defer func() { + <-decodeWorkerSem + }() + return c.decodeWorker( + workerCtx, decodeChunk{items: existingChunk, resumeToken: rtText}, resultChan, req) + } + }(existingChunk)) existingChunk = make([]recordItems, 0, pullRecordsItemsChunkSize) return nil } @@ -616,6 +617,7 @@ func (c *MongoConnector) PullRecords( finishErr := dispatchChunk() // Close the sender chan and wait for events to drain. close(sendChan) + // NB: workerEg cannot be reused after wait() has been called. if err := workerEg.Wait(); err != nil { return err } @@ -638,10 +640,15 @@ func (c *MongoConnector) PullRecords( cancelTimeout() timeoutCtx, cancelTimeout = context.WithTimeout(ctx, time.Hour) + // reset worker timeout. + workerCtxCancel() + workerCtx, workerCtxCancel = context.WithCancel(ctx) + workerEg, workerCtx = errgroup.WithContext(workerCtx) + // Restart the send loop. sendChan = make(chan chan sendChunk, workerBufferedChanSize) workerEg.Go(func() error { - return c.sendLoop(workerCtx, sendChan, req) + return c.sendLoop(workerCtx, sendChan, req, &signalledAsNonEmpty) }) // set resume point based on whether operation time should be used or not From 7ec39fc2f2bf7cf2028febc5e1e624e13d5a8f94 Mon Sep 17 00:00:00 2001 From: Bilal Akhtar Date: Wed, 9 Sep 2026 11:24:27 -0400 Subject: [PATCH 5/8] Abstract out worker pool logic to pkg/common/ --- flow/connectors/mongo/cdc.go | 270 +++++++++++--------------------- flow/pkg/common/workers.go | 169 ++++++++++++++++++++ flow/pkg/common/workers_test.go | 81 ++++++++++ flow/pkg/go.mod | 3 +- 4 files changed, 341 insertions(+), 182 deletions(-) create mode 100644 flow/pkg/common/workers.go create mode 100644 flow/pkg/common/workers_test.go diff --git a/flow/connectors/mongo/cdc.go b/flow/connectors/mongo/cdc.go index 1dd63161b..df636949a 100644 --- a/flow/connectors/mongo/cdc.go +++ b/flow/connectors/mongo/cdc.go @@ -16,7 +16,6 @@ import ( "go.mongodb.org/mongo-driver/v2/mongo/options" "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/trace" - "golang.org/x/sync/errgroup" "github.com/PeerDB-io/peerdb/flow/alerting" "github.com/PeerDB-io/peerdb/flow/generated/protos" @@ -231,74 +230,43 @@ func decodeEvent( return nil } -// Constants used by PullRecords. +// Constant used by PullRecords. // -// Buffered channel size for channels to pass records to decode and send loops (see below). -// This should ideally be larger than decodeWorkerBufSize. -const workerBufferedChanSize = 10 - -// Number of recordItems to pass in one chunk to decode/send loops. Doing per-item channel sends -// results in too much coordination and reduces effective concurrency in practice. +// 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 -type decodeChunk struct { - // base64-encoded resume token - resumeToken string - items []recordItems -} - -type sendChunk struct { - // base64-encoded resume token - resumeToken string - // records to send to RecordStream. - records []model.Record[model.RecordItems] -} - -// PullRecords spins up two kinds of goroutines: one sendLoop that lives for the duration of one -// ChangeStream and is responsible for updating the latest checkpoint info. The other type is implemented -// in decodeWorker and is spun up spun up per-chunk (with parallelism managed by decodeWorkerSem), to -// decode one chunk at a time. The context passed into these goroutines is shared, but is not used +// PullRecords spins up worker goroutines using PullRecordsWorkerPool. +// The context passed into these goroutines is not used to signal timeouts // to signal timeouts; rather, when cutting a batch, the `records` chan passed to sendLoop below is // closed, and we wait for all child goroutines to gracefully drain. -func (c *MongoConnector) sendLoop( +func (c *MongoConnector) recordSender( ctx context.Context, - records <-chan chan sendChunk, + records []model.Record[model.RecordItems], + resumeToken string, req *model.PullRecordsRequest[model.RecordItems], signalledAsNonEmpty *bool, ) error { - for { - select { - case recordChan, ok := <-records: - if !ok { - return nil - } - select { - case sendChunk := <-recordChan: - for i := range sendChunk.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, sendChunk.records[i]); err != nil { - return err - } - } - if sendChunk.resumeToken != "" { - req.RecordStream.UpdateLatestCheckpointText(sendChunk.resumeToken) - } - case <-ctx.Done(): - return nil - } - case <-ctx.Done(): - return nil + 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 recordItems struct { +type encodedMongoEvent struct { maybeFullDocument *bson.Raw operationType operationType sourceTableName string @@ -307,28 +275,26 @@ type recordItems struct { commitTimeNanos int64 } -// decodeWorker is spun up by PullRecords in separate goroutines, up to -// PEERDB_MONGODB_NUM_PARALLEL_DECODE_THREADS in parallel. The output is sent through `send` to -// `sendLoop` directly. -func (c *MongoConnector) decodeWorker( - ctx context.Context, - chunk decodeChunk, - send chan<- sendChunk, +// 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], -) error { +) ([]model.Record[model.RecordItems], error) { // Utils used by this routine. converter := NewDirectBsonConverter() fullDocumentColumnName := DefaultFullDocumentColumnName if req.InternalVersion < shared.InternalVersion_MongoDBFullDocumentColumnToDoc { fullDocumentColumnName = LegacyFullDocumentColumnName } - parseItem := func(item recordItems) (model.Record[model.RecordItems], error) { + parseItem := func(event encodedMongoEvent) (model.Record[model.RecordItems], error) { items := model.NewRecordItems(2) - if len(item.documentKey) > 0 { - rv := item.documentKey.Lookup(DefaultDocumentKeyColumnName) + if len(event.documentKey) > 0 { + rv := event.documentKey.Lookup(DefaultDocumentKeyColumnName) if rv.IsZero() || rv.Type == bson.TypeNull { - return nil, exceptions.NewInvalidIdValueError(item.sourceTableName) + return nil, exceptions.NewInvalidIdValueError(event.sourceTableName) } qValue, err := converter.QValueStringFromId(rv, req.InternalVersion) if err != nil { @@ -339,8 +305,8 @@ func (c *MongoConnector) decodeWorker( return nil, fmt.Errorf("document key is nil") } - if item.maybeFullDocument != nil && len(*item.maybeFullDocument) > 0 { - qValue, err := converter.QValueJSONFromDocument(*item.maybeFullDocument) + 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) } @@ -354,45 +320,40 @@ func (c *MongoConnector) decodeWorker( items.AddColumn(fullDocumentColumnName, types.QValueJSON{Val: "{}"}) } var record model.Record[model.RecordItems] - switch item.operationType { + switch event.operationType { case operationTypeInsert: record = &model.InsertRecord[model.RecordItems]{ - BaseRecord: model.BaseRecord{CommitTimeNano: item.commitTimeNanos}, + BaseRecord: model.BaseRecord{CommitTimeNano: event.commitTimeNanos}, Items: items, - SourceTableName: item.sourceTableName, - DestinationTableName: item.destinationTableName, + SourceTableName: event.sourceTableName, + DestinationTableName: event.destinationTableName, } case operationTypeUpdate, operationTypeReplace: record = &model.UpdateRecord[model.RecordItems]{ - BaseRecord: model.BaseRecord{CommitTimeNano: item.commitTimeNanos}, + BaseRecord: model.BaseRecord{CommitTimeNano: event.commitTimeNanos}, NewItems: items, - SourceTableName: item.sourceTableName, - DestinationTableName: item.destinationTableName, + SourceTableName: event.sourceTableName, + DestinationTableName: event.destinationTableName, } case operationTypeDelete: record = &model.DeleteRecord[model.RecordItems]{ - BaseRecord: model.BaseRecord{CommitTimeNano: item.commitTimeNanos}, + BaseRecord: model.BaseRecord{CommitTimeNano: event.commitTimeNanos}, Items: items, - SourceTableName: item.sourceTableName, - DestinationTableName: item.destinationTableName, + SourceTableName: event.sourceTableName, + DestinationTableName: event.destinationTableName, } } return record, nil } - modelRecords := make([]model.Record[model.RecordItems], len(chunk.items)) - for i := range chunk.items { + modelRecords := make([]model.Record[model.RecordItems], len(events)) + for i := range events { var err error - if modelRecords[i], err = parseItem(chunk.items[i]); err != nil { - return err + if modelRecords[i], err = parseItem(events[i]); err != nil { + return nil, err } } - select { - case send <- sendChunk{records: modelRecords, resumeToken: chunk.resumeToken}: - case <-ctx.Done(): - return ctx.Err() - } - return nil + return modelRecords, nil } func (c *MongoConnector) PullRecords( @@ -500,20 +461,24 @@ func (c *MongoConnector) PullRecords( slog.Int("channelLen", req.RecordStream.ChannelLen()), slog.Float64("elapsedMinutes", time.Since(pullStart).Minutes())) }() - // Context inheritance tree: There's a parent ctx (always called ctx in this function). - // Off of that, we create two child contexts, one for workers (called workerCtx), that's passed - // to children goroutines (eg. decodeWorker). The other, timeoutCtx, is managed by us for timeouts. - // Note that when we hit a timeout, we don't want to cancel the workerCtx; we want the workers - // to gracefully drain. We recreate workerCtx and workerEg anytime we recreate the changestream. - // - // Also note that timeoutCtx is occasionally rewritten, such as when the first record arrives, - // or when we reset the changestream. - workerCtx, workerCtxCancel := context.WithCancel(ctx) + numParallelDecodeWorkers, err := internal.PeerDBMongoDBNumParallelDecodeThreads(ctx, req.Env) + if err != nil { + return err + } + workerPool := common.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) - var workerEg *errgroup.Group - workerEg, workerCtx = errgroup.WithContext(workerCtx) reportBytesShutdown := common.Interval(ctx, time.Second*10, func() { read := deltaBytesProcessed.Swap(0) @@ -523,8 +488,7 @@ func (c *MongoConnector) PullRecords( defer func() { cancelTimeout() - workerCtxCancel() - _ = workerEg.Wait() + _ = workerPool.Wait(ctx) reportBytesShutdown() read := deltaBytesProcessed.Swap(0) otelManager.Metrics.FetchedBytesCounter.Add(ctx, read) @@ -564,65 +528,6 @@ func (c *MongoConnector) PullRecords( slog.Float64("elapsedMinutes", time.Since(pullStart).Minutes())) } } - existingChunk := make([]recordItems, 0, pullRecordsItemsChunkSize) - sendChan := make(chan chan sendChunk, workerBufferedChanSize) - numParallelDecodeWorkers, err := internal.PeerDBMongoDBNumParallelDecodeThreads(ctx, req.Env) - if err != nil { - return err - } - decodeWorkerSem := make(chan struct{}, max(1, numParallelDecodeWorkers)) - workerEg.Go(func() error { - return c.sendLoop(workerCtx, sendChan, req, &signalledAsNonEmpty) - }) - - dispatchChunk := func() error { - if len(existingChunk) == 0 { - // Nothing to send. - return nil - } - 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) - } - select { - case decodeWorkerSem <- struct{}{}: - case <-workerCtx.Done(): - return workerEg.Wait() - } - resultChan := make(chan sendChunk) - select { - case sendChan <- resultChan: - case <-workerCtx.Done(): - return workerEg.Wait() - } - workerEg.Go(func(existingChunk []recordItems) func() error { - // Two-level method to capture existingChunk before it changes - // below. - return func() error { - defer func() { - <-decodeWorkerSem - }() - return c.decodeWorker( - workerCtx, decodeChunk{items: existingChunk, resumeToken: rtText}, resultChan, req) - } - }(existingChunk)) - existingChunk = make([]recordItems, 0, pullRecordsItemsChunkSize) - return nil - } - - drainWorkers := func() error { - finishErr := dispatchChunk() - // Close the sender chan and wait for events to drain. - close(sendChan) - // NB: workerEg cannot be reused after wait() has been called. - if err := workerEg.Wait(); err != nil { - return err - } - return finishErr - } recreateChangeStream := func(useOperationTime bool) error { // extract the most recent resumeToken @@ -640,16 +545,8 @@ func (c *MongoConnector) PullRecords( cancelTimeout() timeoutCtx, cancelTimeout = context.WithTimeout(ctx, time.Hour) - // reset worker timeout. - workerCtxCancel() - workerCtx, workerCtxCancel = context.WithCancel(ctx) - workerEg, workerCtx = errgroup.WithContext(workerCtx) - - // Restart the send loop. - sendChan = make(chan chan sendChunk, workerBufferedChanSize) - workerEg.Go(func() error { - return c.sendLoop(workerCtx, sendChan, req, &signalledAsNonEmpty) - }) + // 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 { @@ -680,7 +577,11 @@ func (c *MongoConnector) PullRecords( return fmt.Errorf("unexpected: changestream.Next() returned false but no change stream error was recorded") } - if err := drainWorkers(); err != nil { + if err := workerPool.Flush(ctx); err != nil { + return err + } + + if err := workerPool.Wait(ctx); err != nil { return err } @@ -752,7 +653,7 @@ func (c *MongoConnector) PullRecords( continue } - items := recordItems{ + event := encodedMongoEvent{ documentKey: changeEvent.DocumentKey, maybeFullDocument: changeEvent.FullDocument, operationType: operationType(changeEvent.OperationType), @@ -760,7 +661,7 @@ func (c *MongoConnector) PullRecords( destinationTableName: destinationTableName, commitTimeNanos: commitTimeNanos, } - switch items.operationType { + switch event.operationType { case operationTypeInsert, operationTypeReplace, operationTypeUpdate, operationTypeDelete: // Happy path. incrementRecordCount() @@ -784,15 +685,22 @@ func (c *MongoConnector) PullRecords( continue } otelManager.Metrics.FetchedEventSizeHistogram.Record(ctx, changeEventSize) - existingChunk = append(existingChunk, items) - if len(existingChunk) >= pullRecordsItemsChunkSize { - if err := dispatchChunk(); err != nil { - return err - } + 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 := drainWorkers(); err != nil { + if err := workerPool.Wait(ctx); err != nil { return err } diff --git a/flow/pkg/common/workers.go b/flow/pkg/common/workers.go new file mode 100644 index 000000000..153219687 --- /dev/null +++ b/flow/pkg/common/workers.go @@ -0,0 +1,169 @@ +package common + +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 decode and send loops (see below). +// This should ideally be larger than decodeWorkerBufSize. +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{}, p.Concurrency) + p.sender = make(chan chan sendMsg[D, RT], workerBufferedChanSize) + 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 nil + } + return nil + }) + p.inProgressChunk = make([]E, 0, p.ChunkSize) + return nil + case <-ctx.Done(): + p.ctxCancel() + return ctx.Err() + case <-p.workerCtx.Done(): + return nil + } +} + +// 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/pkg/common/workers_test.go b/flow/pkg/common/workers_test.go new file mode 100644 index 000000000..626f27602 --- /dev/null +++ b/flow/pkg/common/workers_test.go @@ -0,0 +1,81 @@ +package common + +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()) + + // Keep feeding past the chunk that fails: the pool should wind itself down instead of + // wedging, and AddItem stays quiet because Wait is what reports the worker's error. + for i := range 20 { + require.NoError(t, pool.AddItem(t.Context(), i, fmt.Sprintf("token-%d", i))) + } + require.NoError(t, pool.Flush(t.Context())) + require.ErrorIs(t, pool.Wait(t.Context()), errDecode) + require.Equal(t, "token-3", lastResumeToken) +} diff --git a/flow/pkg/go.mod b/flow/pkg/go.mod index 3dd7d0cb6..51b6c443b 100644 --- a/flow/pkg/go.mod +++ b/flow/pkg/go.mod @@ -13,6 +13,8 @@ require ( github.com/stretchr/testify v1.12.1 go.mongodb.org/mongo-driver/v2 v2.8.0 go.temporal.io/sdk v1.47.0 + go.uber.org/goleak v1.3.0 + golang.org/x/sync v0.22.0 google.golang.org/api v0.287.1 ) @@ -93,7 +95,6 @@ require ( golang.org/x/mod v0.37.0 // indirect golang.org/x/net v0.57.0 // indirect golang.org/x/oauth2 v0.36.0 // indirect - golang.org/x/sync v0.22.0 // indirect golang.org/x/sys v0.47.0 // indirect golang.org/x/telemetry v0.0.0-20260625142307-59b4966ccb57 // indirect golang.org/x/text v0.40.0 // indirect From 552920726f7acc51751881c4b03a91f2371db87b Mon Sep 17 00:00:00 2001 From: Bilal Akhtar Date: Thu, 10 Sep 2026 12:27:26 -0400 Subject: [PATCH 6/8] Claude suggested fixes --- flow/connectors/mongo/cdc.go | 6 +++--- flow/pkg/common/workers.go | 12 ++++++------ 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/flow/connectors/mongo/cdc.go b/flow/connectors/mongo/cdc.go index df636949a..e690422e2 100644 --- a/flow/connectors/mongo/cdc.go +++ b/flow/connectors/mongo/cdc.go @@ -238,9 +238,9 @@ func decodeEvent( const pullRecordsItemsChunkSize = 256 // PullRecords spins up worker goroutines using PullRecordsWorkerPool. -// The context passed into these goroutines is not used to signal timeouts -// to signal timeouts; rather, when cutting a batch, the `records` chan passed to sendLoop below is -// closed, and we wait for all child goroutines to gracefully drain. +// 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], diff --git a/flow/pkg/common/workers.go b/flow/pkg/common/workers.go index 153219687..cf6f12f26 100644 --- a/flow/pkg/common/workers.go +++ b/flow/pkg/common/workers.go @@ -42,8 +42,8 @@ type PullRecordsWorkerPool[E, D, RT any] struct { closed bool } -// Buffered channel size for channels to pass records to decode and send loops (see below). -// This should ideally be larger than decodeWorkerBufSize. +// 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() @@ -51,8 +51,8 @@ const workerBufferedChanSize = 10 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{}, p.Concurrency) - p.sender = make(chan chan sendMsg[D, RT], workerBufferedChanSize) + 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 @@ -130,7 +130,7 @@ func (p *PullRecordsWorkerPool[E, D, RT]) Flush(ctx context.Context) error { select { case sendChan <- sendMsg: case <-p.workerCtx.Done(): - return nil + return context.Cause(p.workerCtx) } return nil }) @@ -140,7 +140,7 @@ func (p *PullRecordsWorkerPool[E, D, RT]) Flush(ctx context.Context) error { p.ctxCancel() return ctx.Err() case <-p.workerCtx.Done(): - return nil + return context.Cause(p.workerCtx) } } From c562b7f77b9846c2c42fab3ae1fe78cb1cc39689 Mon Sep 17 00:00:00 2001 From: Bilal Akhtar Date: Thu, 10 Sep 2026 13:05:40 -0400 Subject: [PATCH 7/8] test fix --- flow/pkg/common/workers_test.go | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/flow/pkg/common/workers_test.go b/flow/pkg/common/workers_test.go index 626f27602..9da5c2556 100644 --- a/flow/pkg/common/workers_test.go +++ b/flow/pkg/common/workers_test.go @@ -70,12 +70,15 @@ func TestPullRecordsWorkerPoolErrorFromWorker(t *testing.T) { } pool.Init(t.Context()) - // Keep feeding past the chunk that fails: the pool should wind itself down instead of - // wedging, and AddItem stays quiet because Wait is what reports the worker's error. + errSeen := false for i := range 20 { - require.NoError(t, pool.AddItem(t.Context(), i, fmt.Sprintf("token-%d", i))) + if err := pool.AddItem(t.Context(), i, fmt.Sprintf("token-%d", i)); err != nil { + require.ErrorIs(t, err, errDecode) + errSeen = true + break + } } - require.NoError(t, pool.Flush(t.Context())) - require.ErrorIs(t, pool.Wait(t.Context()), errDecode) + require.True(t, errSeen, "no error seen from AddItem") + require.Error(t, pool.Wait(t.Context())) require.Equal(t, "token-3", lastResumeToken) } From 366469e316fc5db66913c11e9faa02518b9e7378 Mon Sep 17 00:00:00 2001 From: Bilal Akhtar Date: Thu, 10 Sep 2026 15:23:18 -0400 Subject: [PATCH 8/8] Move worker pool to shared/concurrency/ --- flow/connectors/mongo/cdc.go | 3 ++- flow/go.mod | 4 ++-- flow/pkg/go.mod | 3 +-- flow/{pkg/common => shared/concurrency}/workers.go | 2 +- flow/{pkg/common => shared/concurrency}/workers_test.go | 2 +- 5 files changed, 7 insertions(+), 7 deletions(-) rename flow/{pkg/common => shared/concurrency}/workers.go (99%) rename flow/{pkg/common => shared/concurrency}/workers_test.go (99%) diff --git a/flow/connectors/mongo/cdc.go b/flow/connectors/mongo/cdc.go index e690422e2..6f464a681 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" ) @@ -465,7 +466,7 @@ func (c *MongoConnector) PullRecords( if err != nil { return err } - workerPool := common.PullRecordsWorkerPool[encodedMongoEvent, []model.Record[model.RecordItems], string]{ + workerPool := concurrency.PullRecordsWorkerPool[encodedMongoEvent, []model.Record[model.RecordItems], string]{ Concurrency: int(numParallelDecodeWorkers), ChunkSize: pullRecordsItemsChunkSize, WorkerFunc: func(events []encodedMongoEvent) ([]model.Record[model.RecordItems], error) { diff --git a/flow/go.mod b/flow/go.mod index e0453b7d5..066d0fcdb 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/pkg/go.mod b/flow/pkg/go.mod index 51b6c443b..3dd7d0cb6 100644 --- a/flow/pkg/go.mod +++ b/flow/pkg/go.mod @@ -13,8 +13,6 @@ require ( github.com/stretchr/testify v1.12.1 go.mongodb.org/mongo-driver/v2 v2.8.0 go.temporal.io/sdk v1.47.0 - go.uber.org/goleak v1.3.0 - golang.org/x/sync v0.22.0 google.golang.org/api v0.287.1 ) @@ -95,6 +93,7 @@ require ( golang.org/x/mod v0.37.0 // indirect golang.org/x/net v0.57.0 // indirect golang.org/x/oauth2 v0.36.0 // indirect + golang.org/x/sync v0.22.0 // indirect golang.org/x/sys v0.47.0 // indirect golang.org/x/telemetry v0.0.0-20260625142307-59b4966ccb57 // indirect golang.org/x/text v0.40.0 // indirect diff --git a/flow/pkg/common/workers.go b/flow/shared/concurrency/workers.go similarity index 99% rename from flow/pkg/common/workers.go rename to flow/shared/concurrency/workers.go index cf6f12f26..d7d60c322 100644 --- a/flow/pkg/common/workers.go +++ b/flow/shared/concurrency/workers.go @@ -1,4 +1,4 @@ -package common +package concurrency import ( "context" diff --git a/flow/pkg/common/workers_test.go b/flow/shared/concurrency/workers_test.go similarity index 99% rename from flow/pkg/common/workers_test.go rename to flow/shared/concurrency/workers_test.go index 9da5c2556..64ee7fdb9 100644 --- a/flow/pkg/common/workers_test.go +++ b/flow/shared/concurrency/workers_test.go @@ -1,4 +1,4 @@ -package common +package concurrency import ( "context"