jiangzzhu opened a new pull request, #40870:
URL: https://github.com/apache/beam/pull/40870
## Summary
Fixes a partition stranding issue in
`DetectNewPartitionsAction.schedulePartitions` when `DetectNewPartitionsDoFn`'s
`TimestampRange` restriction is dynamically split across a multi-timestamp
batch:
- **Problem**: `DetectNewPartitionsAction.schedulePartitions` groups new
partitions by `CreatedAt` in a `TreeMap` and accumulates groups across multiple
`CreatedAt` timestamps (`T1 < T2 < ... < Tk`) up to 200 partitions before
flushing. Previously, when flushing a batch at `Tk`, it:
1. Executed `updateBatchToScheduled(batchPartitionsDifferentCreatedAt)`
(committing `State = SCHEDULED` in Spanner for all partitions across `T1..Tk`)
**before** calling `tracker.tryClaim(Tk)`.
2. Only called `tracker.tryClaim` on the batch's largest timestamp `Tk`
(never claiming `T1..T_{k-1}`).
3. If the runner dynamically split the restriction at `T_split` (`T1 <
T_split <= Tk`), `tracker.tryClaim(Tk)` returned `false` and
`schedulePartitions` immediately returned `ProcessContinuation.stop()` without
calling `outputBatch(...)`. The primary restriction never emitted the `T1`
partitions, while the residual restriction `[T_split, MAX)` only queries
`CreatedAt > T_split`, permanently stranding the earlier `CreatedAt < T_split`
partitions in `State = SCHEDULED` with no reader and freezing the pipeline
watermark.
- **Fix**:
- Call `tracker.tryClaim(batchCreatedAt)` for each `CreatedAt` group (in
ascending `CreatedAt` order) **before** adding that group's partitions to
`batchPartitionsDifferentCreatedAt` or mutating Spanner metadata state.
- If `tracker.tryClaim(batchCreatedAt)` returns `false`, flush
(`updateBatchToScheduled` and `outputBatch`) any already-claimed groups
accumulated in `batchPartitionsDifferentCreatedAt` before returning
`ProcessContinuation.stop()`, leaving the unclaimed group in `State = CREATED`
for the residual restriction.
- Preserve the 200-partition Spanner write batching across claimed
`CreatedAt` groups.
## Testing
- Added unit tests in `DetectNewPartitionsActionTest` verifying:
- `tryClaim` is invoked before `dao.updateToScheduled`.
- Stopping without calling `dao.updateToScheduled` or
`receiver.outputWithTimestamp` when the first `CreatedAt` group cannot be
claimed.
- Flushing (`updateToScheduled` + `outputWithTimestamp`) only the
already-claimed `CreatedAt` groups when a subsequent group cannot be claimed
due to a restriction split.
- Flushing both a full 200-partition batch and a subsequent partial
claimed batch before stopping on an unclaimed group.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
