[Dataflow Streaming] [Multi Key] Introduce KeyGroupWorkQueue and integrate with BoundedQueueExecutor#38767
[Dataflow Streaming] [Multi Key] Introduce KeyGroupWorkQueue and integrate with BoundedQueueExecutor#38767arunpandianp wants to merge 1 commit into
Conversation
…oundedQueueExecutor
- Add Uint128Proto and key_group to WorkItem in windmill.proto.
- Update Work model to support key groups.
- Implement KeyGroupWorkQueue for grouping tasks by key group. KepGroupWorkQueue
is a FIFO queue that allows polling elements based on
global order and also by order within a key group.
- Integrate BoundedQueueExecutor with KepGroupWorkQueue and support targeted polling.
|
R: @scwhittle |
|
Stopping reviewer notifications for this pull request: review requested by someone other than the bot, ceding control. If you'd like to restart, comment |
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request introduces a new KeyGroupWorkQueue implementation to support grouping and targeted polling of work items by key group. It updates the Work model and windmill.proto to include key group information and integrates this new queue mechanism into the BoundedQueueExecutor. These changes are gated behind an experimental flag, ensuring no impact on existing core logic. Highlights
New Features🧠 You can now enable Memory (public preview) to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for multi-key bundles in the Streaming Dataflow Worker, guarded by the unstable_enable_multi_key_bundle experiment. It adds a custom, thread-safe KeyGroupWorkQueue to allow polling work by computation and key group in FIFO order, along with corresponding updates to BoundedQueueExecutor, Work, and windmill.proto. Feedback on the new queue implementation suggests adding a defensive check in unlinkNode to prevent double-unlinking issues, and reassigning the result of checkStateNotNull in take() to satisfy static analysis nullness checks.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| private void unlinkNode(Node node) { | ||
| // 1. Unlink from global list | ||
| Node prevG = node.prevNode; | ||
| Node nextG = node.nextNode; | ||
| if (prevG != null && nextG != null) { | ||
| prevG.nextNode = nextG; | ||
| nextG.prevNode = prevG; | ||
| } | ||
| node.prevNode = null; | ||
| node.nextNode = null; |
There was a problem hiding this comment.
To prevent potential issues with double-unlinking (e.g., if unlinkNode is called multiple times on the same node due to future refactoring or concurrent edge cases), we should add a defensive check at the beginning of the method. If both prevNode and nextNode are null, the node is already unlinked, and we can return early to avoid corrupting the queue's size and internal state.
private void unlinkNode(Node node) {
if (node.prevNode == null && node.nextNode == null) {
return;
}
// 1. Unlink from global list
Node prevG = node.prevNode;
Node nextG = node.nextNode;
if (prevG != null && nextG != null) {
prevG.nextNode = nextG;
nextG.prevNode = prevG;
}
node.prevNode = null;
node.nextNode = null;| @Nullable Node node = removeFirstGlobal(); | ||
| checkStateNotNull(node, "Queue is empty but size was " + size); | ||
| Runnable task = node.task; | ||
| checkStateNotNull(task, "Encountered null task in queue"); | ||
| return task; |
There was a problem hiding this comment.
In environments with strict nullness checking (such as the Checker Framework used in Beam), calling checkStateNotNull without reassigning the result to a non-nullable variable can still trigger static analysis warnings when dereferencing the original @Nullable reference. Reassigning the non-null result of checkStateNotNull to a non-nullable variable is the standard and safe way to satisfy the compiler and static analysis tools.
| @Nullable Node node = removeFirstGlobal(); | |
| checkStateNotNull(node, "Queue is empty but size was " + size); | |
| Runnable task = node.task; | |
| checkStateNotNull(task, "Encountered null task in queue"); | |
| return task; | |
| Node node = checkStateNotNull(removeFirstGlobal(), "Queue is empty but size was " + size); | |
| Runnable task = checkStateNotNull(node.task, "Encountered null task in queue"); | |
| return task; |
1 participant
The changes are behind an experiment
unstable_enable_multi_key_bundleand does not affect the default core logic.