fix(rust): finish query streams after a zero-duration timeout (#4506)

A positive query timeout emits one error and then ends, but
`Duration::ZERO` leaves `TimeoutStream` in `NotStarted`, returning the
same error on every subsequent poll.

Transition the zero-duration path to `Completed` before returning its
error. The initial error is unchanged. Extend both timeout tests to
check that subsequent polls return `None`; the normal-completion test is
unchanged. This affects callers that continue polling after an error,
not `try_collect`, which already stops at the first error.

Validation:
- `cargo fmt --all` and `cargo fmt --all --check`: pass.
- Mechanically extracted the unchanged implementation and three upstream
timeout tests into a native Rust test binary using real cached Arrow
58.4.0, DataFusion 54.1.0, futures 0.3.32 and Tokio 1.53.1 libraries
(matching the current lock). With the new assertions, baseline: 2 pass /
1 fail; fix: 3 pass.
- Full-crate `cargo test --locked --offline --features remote -p lancedb
--lib test_timeout_stream` could not start because this environment
lacks the newly locked Lance `v14.0.0-beta.10` checkout. The component
comparison is not a claim that the full Cargo suite passed; CI still
needs to validate integration.

Signed-off-by: Xucheng Zhou <aden1350@outlook.com>
This commit is contained in:
adenzhou1350
2026-10-23 03:14:45 +08:00
committed by GitHub
parent 3d883af674
commit 93f054bcc5
+7 -1
View File
@@ -686,7 +686,9 @@ impl Stream for TimeoutStream {
match &mut self.state {
TimeoutState::NotStarted { timeout } => {
if timeout.is_zero() {
return std::task::Poll::Ready(Some(Err(Self::timeout_error(timeout))));
let err = Self::timeout_error(timeout);
self.state = TimeoutState::Completed;
return std::task::Poll::Ready(Some(Err(err)));
}
let deadline = Box::pin(tokio::time::sleep(*timeout));
self.state = TimeoutState::Started {
@@ -1159,6 +1161,8 @@ mod tests {
.to_string()
.contains("Query timeout")
);
assert!(timeout_stream.next().await.is_none());
assert!(timeout_stream.next().await.is_none());
}
#[tokio::test]
@@ -1178,6 +1182,8 @@ mod tests {
let result = timeout_stream.next().await.unwrap();
assert!(result.is_err());
assert!(result.unwrap_err().to_string().contains("Query timeout"));
assert!(timeout_stream.next().await.is_none());
assert!(timeout_stream.next().await.is_none());
}
#[tokio::test]