From 42d6187cde0e6e862505d25d5e5f247b760584dd Mon Sep 17 00:00:00 2001 From: ayu Date: Wed, 29 Jul 2026 23:37:43 +0000 Subject: [PATCH] =?UTF-8?q?=F0=9F=A4=96=20retry=20transient=20snowflake=20?= =?UTF-8?q?errors?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- osprey_coordinator/src/snowflake_client.rs | 179 +++++++++++++++++++++ 1 file changed, 179 insertions(+) diff --git a/osprey_coordinator/src/snowflake_client.rs b/osprey_coordinator/src/snowflake_client.rs index 174c161b..bda3017b 100644 --- a/osprey_coordinator/src/snowflake_client.rs +++ b/osprey_coordinator/src/snowflake_client.rs @@ -1,5 +1,11 @@ +use reqwest::StatusCode; use serde::{Deserialize, Serialize}; +use std::time::Duration; use thiserror::Error; +use tokio::time::sleep; + +const MAX_RETRIES: usize = 3; +const RETRY_BASE_DELAY_MILLIS: u64 = 50; #[derive(Debug, Serialize, Deserialize)] struct SnowflakeRequest { @@ -15,6 +21,22 @@ pub enum SnowflakeClientError { NoIdGeneratedError, } +impl SnowflakeClientError { + fn is_retryable(&self) -> bool { + let Self::IdGenerateError(error) = self else { + return false; + }; + if let Some(status) = error.status() { + return status == StatusCode::TOO_MANY_REQUESTS || status.is_server_error(); + } + error.is_connect() || error.is_timeout() || error.is_body() || error.is_request() + } +} + +fn retry_delay(attempt: usize) -> Duration { + Duration::from_millis((attempt + 1).pow(2) as u64 * RETRY_BASE_DELAY_MILLIS) +} + pub struct SnowflakeClient { snowflake_api_endpoint: String, reqwest_client: reqwest::Client, @@ -41,6 +63,19 @@ impl SnowflakeClient { } pub async fn generate_id(&self) -> Result { + for attempt in 0..=MAX_RETRIES { + match self.generate_id_once().await { + Ok(id) => return Ok(id), + Err(error) if attempt < MAX_RETRIES && error.is_retryable() => { + sleep(retry_delay(attempt)).await; + } + Err(error) => return Err(error), + } + } + unreachable!() + } + + async fn generate_id_once(&self) -> Result { // osprey-snowflake api spec: https://github.com/ayubun/snowflake-id-worker?tab=readme-ov-file#api-spec let mut snowflake_response: Vec = self .reqwest_client @@ -48,6 +83,7 @@ impl SnowflakeClient { .json(&SnowflakeRequest { count: 1 }) .send() .await? + .error_for_status()? .json() .await?; @@ -102,3 +138,146 @@ impl SnowflakeClient { // self.last_buffer_fill = SystemTime::now(); // } // } + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, + }; + use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::TcpListener, + task::JoinHandle, + }; + + enum TestResponse { + Disconnect, + Status(StatusCode, &'static str), + } + + struct TestServer { + endpoint: String, + requests: Arc, + handle: JoinHandle<()>, + } + + impl TestServer { + async fn start(responses: Vec) -> Self { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let requests = Arc::new(AtomicUsize::new(0)); + let request_counter = Arc::clone(&requests); + let handle = tokio::spawn(async move { + for response in responses { + let (mut stream, _) = listener.accept().await.unwrap(); + request_counter.fetch_add(1, Ordering::SeqCst); + match response { + TestResponse::Disconnect => continue, + TestResponse::Status(status, body) => { + let mut request = [0; 1024]; + let _ = stream.read(&mut request).await; + let response = format!( + "HTTP/1.1 {} {}\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + status.as_u16(), + status.canonical_reason().unwrap_or("unknown"), + body.len(), + body + ); + stream.write_all(response.as_bytes()).await.unwrap(); + } + } + } + }); + Self { + endpoint: format!("http://{address}"), + requests, + handle, + } + } + + fn request_count(&self) -> usize { + self.requests.load(Ordering::SeqCst) + } + } + + impl Drop for TestServer { + fn drop(&mut self) { + self.handle.abort(); + } + } + + #[tokio::test] + async fn retries_rate_limit_response() { + let server = TestServer::start(vec![ + TestResponse::Status(StatusCode::TOO_MANY_REQUESTS, "queue full"), + TestResponse::Status(StatusCode::OK, "[123]"), + ]) + .await; + let client = SnowflakeClient::new(server.endpoint.clone()); + + assert_eq!(client.generate_id().await.unwrap(), 123); + assert_eq!(server.request_count(), 2); + } + + #[tokio::test] + async fn retries_server_error_response() { + let server = TestServer::start(vec![ + TestResponse::Status(StatusCode::SERVICE_UNAVAILABLE, "unavailable"), + TestResponse::Status(StatusCode::OK, "[123]"), + ]) + .await; + let client = SnowflakeClient::new(server.endpoint.clone()); + + assert_eq!(client.generate_id().await.unwrap(), 123); + assert_eq!(server.request_count(), 2); + } + + #[tokio::test] + async fn retries_transport_error() { + let server = TestServer::start(vec![ + TestResponse::Disconnect, + TestResponse::Status(StatusCode::OK, "[123]"), + ]) + .await; + let client = SnowflakeClient::new(server.endpoint.clone()); + + assert_eq!(client.generate_id().await.unwrap(), 123); + assert_eq!(server.request_count(), 2); + } + + #[tokio::test] + async fn does_not_retry_client_error_response() { + let server = TestServer::start(vec![ + TestResponse::Status(StatusCode::BAD_REQUEST, "bad request"), + TestResponse::Status(StatusCode::OK, "[123]"), + ]) + .await; + let client = SnowflakeClient::new(server.endpoint.clone()); + + let error = client.generate_id().await.unwrap_err(); + match error { + SnowflakeClientError::IdGenerateError(source) => { + assert_eq!(source.status(), Some(StatusCode::BAD_REQUEST)); + } + SnowflakeClientError::NoIdGeneratedError => panic!("expected request error"), + } + assert_eq!(server.request_count(), 1); + } + + #[tokio::test] + async fn stops_after_four_transient_attempts() { + let server = TestServer::start(vec![ + TestResponse::Status(StatusCode::SERVICE_UNAVAILABLE, "unavailable"), + TestResponse::Status(StatusCode::SERVICE_UNAVAILABLE, "unavailable"), + TestResponse::Status(StatusCode::SERVICE_UNAVAILABLE, "unavailable"), + TestResponse::Status(StatusCode::SERVICE_UNAVAILABLE, "unavailable"), + ]) + .await; + let client = SnowflakeClient::new(server.endpoint.clone()); + + assert!(client.generate_id().await.is_err()); + assert_eq!(server.request_count(), 4); + } +}