diff --git a/rust/lancedb/src/remote/job.rs b/rust/lancedb/src/remote/job.rs index 41d8971ac..7d2e70ac9 100644 --- a/rust/lancedb/src/remote/job.rs +++ b/rust/lancedb/src/remote/job.rs @@ -18,7 +18,9 @@ use crate::remote::client::{HttpSend, RequestResultExt, RestfulLanceDbClient}; /// Delay before the second job-state poll; doubles up to [`MAX_POLL_INTERVAL`]. const INITIAL_POLL_INTERVAL: Duration = Duration::from_millis(200); -const MAX_POLL_INTERVAL: Duration = Duration::from_secs(5); +/// A job's completion is seen at most this long after it happens. A 5 s cap +/// added 2.5 s on average to every short refresh, as much as the job itself. +const MAX_POLL_INTERVAL: Duration = Duration::from_secs(1); #[derive(Debug, Clone, PartialEq, Eq)] enum JobState { @@ -268,13 +270,18 @@ impl JobHandle for RemoteJob { #[cfg(test)] mod tests { + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::time::Duration; + use async_trait::async_trait; use crate::Result; use crate::function::{FunctionVersion, RefreshColumnResult}; use crate::job::{Job, JobHandle, TerminalResult}; + use crate::remote::client::test_utils::client_with_handler; - use super::DescribeJobResponse; + use super::{DescribeJobResponse, RemoteJob}; const FUNCTION_JOB: &str = include_str!("../../tests/fixtures/first_class_functions/v1/remote_function_job.json"); @@ -333,6 +340,31 @@ mod tests { ); } + /// Five in-progress polls then done: 0.2 + 0.4 + 0.8 + 1 + 1 s of waiting. + /// Under the old 5 s cap the same job would have waited 10 s. + #[tokio::test(start_paused = true)] + async fn wait_polls_at_most_every_second() { + let describes = Arc::new(AtomicUsize::new(0)); + let seen = describes.clone(); + let client = client_with_handler(move |request| { + assert_eq!(request.url().path(), "/v1/jobs/describe"); + let n = seen.fetch_add(1, Ordering::SeqCst); + let body = if n < 5 { + r#"{"job_state":"IN_PROGRESS"}"#.to_string() + } else { + UNIT_JOB.to_string() + }; + http::Response::builder().status(200).body(body).unwrap() + }); + let started = tokio::time::Instant::now(); + RemoteJob::new(client, "job-1".to_string()) + .wait() + .await + .expect("unit job completes"); + assert_eq!(describes.load(Ordering::SeqCst), 6); + assert_eq!(started.elapsed(), Duration::from_millis(3400)); + } + #[test] fn remote_wire_unknown_fields_are_forward_decodable() { let response: DescribeJobResponse =