Merge pull request #181 from estie-inc/perf/async-poll-step-up-backoff

perf: replace fixed async poll interval with step-up backoff

Author: yuly3

src/statement/client.rs

@@ -83,15 +83,14 @@ impl StatementApiClient {
         timeout: Duration,
         query_id: Arc<str>,
     ) -> QueryScopedResult<SnowflakeResponse> {
-        const POLL_INTERVAL: Duration = Duration::from_secs(10);
-
         let poll_url = match resolve_poll_url(&self.shared.base_url, poll_relative_url) {
             Ok(url) => url,
             Err(err) => {
                 return Err(QueryScopedError::new(query_id, err));
             }
         };
         let deadline = Instant::now() + timeout;
+        let mut backoff = PollBackoff::new();
 
         loop {
             let resp = self
@@ -151,14 +150,45 @@ impl StatementApiClient {
                     if remaining.is_zero() {
                         return Err(QueryScopedError::new(query_id, TimeoutError::query()));
                     }
-                    sleep(remaining.min(POLL_INTERVAL)).await;
+                    sleep(remaining.min(backoff.next_delay())).await;
                 }
                 Err(err) => return Err(QueryScopedError::new(query_id, err)),
             }
         }
     }
 }
 
+/// Client-side backoff between consecutive poll GETs for an async query.
+///
+/// The poll endpoint is a server-side long-poll: each GET blocks until the query completes or the server-side
+/// timeout elapses, so the client normally does not sleep at all. This backoff only takes effect when a GET
+/// returns in-progress quickly, bounding how tightly the poll loop can spin in that case.
+struct PollBackoff {
+    step: usize,
+}
+
+impl PollBackoff {
+    const STEPS: [Duration; 7] = [
+        Duration::from_millis(500),
+        Duration::from_millis(500),
+        Duration::from_secs(1),
+        Duration::from_millis(1500),
+        Duration::from_secs(2),
+        Duration::from_secs(4),
+        Duration::from_secs(5),
+    ];
+
+    fn new() -> Self {
+        Self { step: 0 }
+    }
+
+    fn next_delay(&mut self) -> Duration {
+        let delay = Self::STEPS[self.step];
+        self.step = (self.step + 1).min(Self::STEPS.len() - 1);
+        delay
+    }
+}
+
 fn resolve_poll_url(
     base_url: &Url,
     poll_relative_url: &str,
@@ -308,6 +338,28 @@ mod tests {
         let _ = server.await;
     }
 
+    #[test]
+    fn poll_backoff_steps_up_then_plateaus_at_five_seconds() {
+        let mut backoff = PollBackoff::new();
+        let observed: Vec<Duration> = (0..9).map(|_| backoff.next_delay()).collect();
+
+        assert_eq!(
+            observed,
+            vec![
+                Duration::from_millis(500),
+                Duration::from_millis(500),
+                Duration::from_secs(1),
+                Duration::from_millis(1500),
+                Duration::from_secs(2),
+                Duration::from_secs(4),
+                Duration::from_secs(5),
+                // Plateaus at the 5s tail once the table is exhausted.
+                Duration::from_secs(5),
+                Duration::from_secs(5),
+            ]
+        );
+    }
+
     #[test]
     fn validate_wire_rejects_empty_named_bind_keys() {
         let parts = into_statement_parts(Statement::new("SELECT :id").bind_named("", 1_i64));

tests/cases/test_async.rs

@@ -7,19 +7,19 @@ use snowflake_connector_rs::{QueryConfig, Result};
 #[tokio::test]
 async fn test_async_query_picked_up_in_final_polling_window() -> Result<()> {
     let query_config =
-        QueryConfig::default().with_async_query_completion_timeout(Duration::from_secs(25));
+        QueryConfig::default().with_async_query_completion_timeout(Duration::from_secs(15));
     let client = common::connect_with_query(query_config)?;
     let session = client.create_session().await?;
 
     let table = session
-        .query(r#"CALL SYSTEM$WAIT(65)"#)
+        .query(r#"CALL SYSTEM$WAIT(50)"#)
         .await?
         .collect_table()
         .await?;
     assert_eq!(table.row_count(), 1);
 
     let value = table.rows::<(String,)>()?.next().unwrap()?.0;
-    assert_eq!(value, "waited 65 seconds");
+    assert_eq!(value, "waited 50 seconds");
 
     Ok(())
 }