test: reproduce metrics loss on early stream termination #341
+68
−0
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
gferrate
force-pushed
the
gferrate/investigate-187-limit-metrics
branch
from
9df7855 to
aa94f31
Compare
gferrate
changed the title
gferrate
commented
Feb 9, 2026
Comment on lines
+395
to
+399
| panic!( | ||
| "Missing metrics for stage key {expected_stage_key:?}. \ | ||
| The LIMIT caused the stream to be dropped before the worker \ | ||
| sent the last FlightData message with metrics." | ||
| ) |
Contributor
Author
There was a problem hiding this comment.
ℹ️ This is the bug. Uncomment #[ignore] to reproduce.
gabotechs
approved these changes
Feb 9, 2026
| /// client drops the stream early. Un-ignore it to reproduce the issue or to verify a fix. | ||
| #[tokio::test] | ||
| #[ignore] | ||
| async fn test_metrics_collection_with_limit_causing_early_stream_termination() { |
Collaborator
There was a problem hiding this comment.
I've seen that another consequence of this is that calling rewrite_distributed_plan_with_metrics on the plan fails with an error that says something like "not enough metrics provided to rewrite task".
This is expected an another way of reproducing the issue. I wonder if rewrite_distributed_plan_with_metrics should be error proof instead?
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Adds a test reproducing the issue described in #187.
Since the test fails by design, it is marked as
#[ignore].The underlying issue is that metrics are attached to the last
FlightDatamessage of the last partition stream on the worker side (indo_get.rs). When aLIMITis present, the client-side network boundary node (e.g.NetworkShuffleExec) drops the stream once it has enough rows, before the worker finishes all partitions. The final message carrying the metrics is never sent or never consumed, so metrics for that stage are lost.The test runs a
GROUP BY ... LIMIT 1query against theflights_1mdataset (1M rows). The large dataset ensures the worker is still producing data when the LIMIT causes the client to drop the stream.Note: This don't actually solve the issue, just showcases that it is really an issue!