From 93f054bcc5b02927720cdcd52fd7b1cc13c9a731 Mon Sep 17 00:00:00 2001
From: adenzhou1350 <aden1350@outlook.com>
Date: Fri, 23 Oct 2026 03:14:45 +0800
Subject: [PATCH] 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>
---
 rust/lancedb/src/utils/mod.rs | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/rust/lancedb/src/utils/mod.rs b/rust/lancedb/src/utils/mod.rs
index 20f209f2e..ebf3f697b 100644
--- a/rust/lancedb/src/utils/mod.rs
+++ b/rust/lancedb/src/utils/mod.rs
@@ -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]