Refactor FSM backoff handling: replace fetch attempts with scheduled timer for retries

Author: fabian4

crates/pavis/src/admin.rs

@@ -7,10 +7,33 @@ use std::sync::Arc;
 use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
 use tokio::net::TcpListener;
 use tokio::sync::watch;
-use tokio::time::Instant;
+use tokio::time::{Duration, Instant, timeout};
 
 use crate::state::RuntimeStateHandle;
 
+const ADMIN_REQUEST_LINE_LIMIT_BYTES: usize = 4096;
+const ADMIN_READ_TIMEOUT: Duration = Duration::from_secs(5);
+
+async fn read_request_line(
+    reader: &mut BufReader<tokio::net::TcpStream>,
+) -> std::io::Result<Option<String>> {
+    let mut buf = Vec::new();
+    let read = timeout(ADMIN_READ_TIMEOUT, reader.read_until(b'\n', &mut buf))
+        .await
+        .map_err(|_| std::io::Error::new(std::io::ErrorKind::TimedOut, "admin read timeout"))??;
+    if read == 0 {
+        return Ok(None);
+    }
+    if buf.len() > ADMIN_REQUEST_LINE_LIMIT_BYTES {
+        return Err(std::io::Error::new(
+            std::io::ErrorKind::InvalidData,
+            "admin request line too long",
+        ));
+    }
+    let line = String::from_utf8_lossy(&buf);
+    Ok(Some(line.trim_end_matches(['\r', '\n']).to_string()))
+}
+
 /// Admin API worker service.
 ///
 /// Provides read-only operational endpoints:
@@ -115,11 +138,9 @@ impl AdminApiWorker {
                     match accept_result {
                         Ok((stream, _peer_addr)) => {
                             let mut reader = BufReader::new(stream);
-                            let mut request_line = String::new();
-
-                            match reader.read_line(&mut request_line).await {
-                                Ok(0) => continue, // Connection closed
-                                Ok(_) => {
+                            match read_request_line(&mut reader).await {
+                                Ok(None) => continue, // Connection closed
+                                Ok(Some(request_line)) => {
                                     // Parse request line: "GET /path HTTP/1.1"
                                     let parts: Vec<&str> = request_line.split_whitespace().collect();
                                     if parts.len() >= 2 {

crates/pavis/src/agent/fsm.rs

@@ -391,16 +391,10 @@ impl Fsm {
             }
             (State::Verifying(_), Event::VerifyFail { etag, now }) => {
                 self.ctx.set_rejected(etag, now);
-                self.ctx.backoff_attempt = 0;
-                self.state = State::Fetching;
-                if let Some(etag) = self.ctx.conditional_etag(now) {
-                    effects.push(Effect::FetchConditional {
-                        wait_ms: WAIT_MS,
-                        etag,
-                    });
-                } else {
-                    effects.push(Effect::FetchUnconditional { wait_ms: WAIT_MS });
-                }
+                let delay = backoff_delay(self.ctx.backoff_attempt);
+                self.ctx.backoff_attempt = self.ctx.backoff_attempt.saturating_add(1);
+                self.state = State::BackoffSleeping;
+                effects.push(Effect::ScheduleTimer { duration: delay });
             }
             (State::Verifying(_), Event::Shutdown) => {
                 self.state = State::Stopped;
@@ -428,16 +422,10 @@ impl Fsm {
             }
             (State::Applying(_), Event::ApplyFail { etag, now }) => {
                 self.ctx.set_rejected(etag, now);
-                self.ctx.backoff_attempt = 0;
-                self.state = State::Fetching;
-                if let Some(etag) = self.ctx.conditional_etag(now) {
-                    effects.push(Effect::FetchConditional {
-                        wait_ms: WAIT_MS,
-                        etag,
-                    });
-                } else {
-                    effects.push(Effect::FetchUnconditional { wait_ms: WAIT_MS });
-                }
+                let delay = backoff_delay(self.ctx.backoff_attempt);
+                self.ctx.backoff_attempt = self.ctx.backoff_attempt.saturating_add(1);
+                self.state = State::BackoffSleeping;
+                effects.push(Effect::ScheduleTimer { duration: delay });
             }
             (State::Applying(_), Event::Shutdown) => {
                 self.state = State::Stopped;
@@ -632,10 +620,8 @@ mod tests {
             fsm.context().last_rejected_etag.as_deref(),
             Some("sha256:bad")
         );
-        assert!(effects.iter().any(|effect| matches!(
-            effect,
-            Effect::FetchConditional { .. } | Effect::FetchUnconditional { .. }
-        )));
+        assert!(matches!(fsm.state, State::BackoffSleeping));
+        assert!(matches!(effects.as_slice(), [Effect::ScheduleTimer { .. }]));
     }
 
     #[test]
@@ -656,10 +642,8 @@ mod tests {
             fsm.context().last_rejected_etag.as_deref(),
             Some("sha256:bad")
         );
-        assert!(effects.iter().any(|effect| matches!(
-            effect,
-            Effect::FetchConditional { .. } | Effect::FetchUnconditional { .. }
-        )));
+        assert!(matches!(fsm.state, State::BackoffSleeping));
+        assert!(matches!(effects.as_slice(), [Effect::ScheduleTimer { .. }]));
     }
 
     #[test]

crates/pavis/src/proxy/service/io.rs

@@ -12,7 +12,10 @@ use pingora::http::ResponseHeader;
 use pingora::prelude::*;
 use pingora::protocols::Digest;
 use pingora::proxy::{ProxyHttp, Session};
+use std::collections::HashSet;
 use std::sync::Arc;
+use std::sync::Mutex;
+use std::sync::OnceLock;
 use std::sync::atomic::Ordering;
 use std::time::Duration;
 use tracing_opentelemetry::OpenTelemetrySpanExt;
@@ -27,6 +30,20 @@ use super::state::Proxy;
 static POOL_KEY_LOG_COUNTER: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
 static SNI_FRAGMENT_WARN_COUNTER: std::sync::atomic::AtomicU64 =
     std::sync::atomic::AtomicU64::new(0);
+static LOGGED_UPSTREAM_CHECK_COUNTER: std::sync::atomic::AtomicU64 =
+    std::sync::atomic::AtomicU64::new(0);
+static LOGGED_UPSTREAM_KEYS: OnceLock<Mutex<HashSet<u64>>> = OnceLock::new();
+
+fn should_log_upstream_config(hash: u64) -> bool {
+    if LOGGED_UPSTREAM_CHECK_COUNTER.fetch_add(1, Ordering::Relaxed) % 256 != 0 {
+        return false;
+    }
+    let set = LOGGED_UPSTREAM_KEYS.get_or_init(|| Mutex::new(HashSet::new()));
+    if let Ok(mut guard) = set.try_lock() {
+        return guard.insert(hash);
+    }
+    false
+}
 
 #[async_trait]
 impl ProxyHttp for Proxy {
@@ -363,15 +380,11 @@ impl ProxyHttp for Proxy {
                         .map(|d| d.as_millis() as u64)
                         .unwrap_or(60000);
 
-                    route_phase.set_retry_context(RetryContext::new(
-                        retry_policy.clone(),
-                        request_timeout_ms,
-                        self.telemetry.metrics.clone(),
-                        upstream_name.clone(),
-                    ));
-
                     let mut body_bytes = Vec::new();
                     let limit = *max_request_body_buffer_bytes;
+                    let mut budget_guard = crate::retry::BudgetGuard::new();
+                    let mut budget_exhausted = false;
+                    let mut truncated = false;
 
                     while let Some(chunk) = session.read_request_body().await? {
                         if (body_bytes.len() as u64) + (chunk.len() as u64) > limit {
@@ -380,19 +393,39 @@ impl ProxyHttp for Proxy {
                                 upstream = %upstream_name,
                                 "Request body exceeds buffer limit, marking as non-replayable"
                             );
-                            body_bytes.extend_from_slice(&chunk);
+                            truncated = true;
+                            break;
+                        }
+                        if !budget_guard.try_add(chunk.len() as u64) {
+                            budget_exhausted = true;
                             break;
                         }
                         body_bytes.extend_from_slice(&chunk);
                     }
 
-                    let buffered = crate::retry::BufferedBody::new(
-                        body_bytes,
-                        limit,
-                        self.telemetry.metrics.clone(),
-                        &upstream_name,
-                    );
-                    route_phase.set_buffered_body(buffered);
+                    if budget_exhausted {
+                        tracing::warn!(
+                            upstream = %upstream_name,
+                            "Retry body buffer budget exhausted; disabling retries for request"
+                        );
+                    } else {
+                        route_phase.set_retry_context(RetryContext::new(
+                            retry_policy.clone(),
+                            request_timeout_ms,
+                            self.telemetry.metrics.clone(),
+                            upstream_name.clone(),
+                        ));
+
+                        let buffered = crate::retry::BufferedBody::new(
+                            body_bytes,
+                            limit,
+                            self.telemetry.metrics.clone(),
+                            &upstream_name,
+                            truncated,
+                            Some(budget_guard),
+                        );
+                        route_phase.set_buffered_body(buffered);
+                    }
                 }
 
                 if let RouteAction::Forward(destinations) = &route.action {
@@ -996,8 +1029,7 @@ impl<'a> UpstreamPeerBuilder<'a> {
             peer.options.tcp_recv_buf = Some(buffer_size as usize);
         }
 
-        // Log effective upstream configuration (once per upstream, using atomic counter to avoid Mutex)
-        static LOGGED_COUNTER: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
+        // Log effective upstream configuration (rate-limited with a sampled, per-upstream set)
         let log_key = format!("{}:{}", upstream_name.0, cluster as *const _ as usize);
         let hash = {
             use std::collections::hash_map::DefaultHasher;
@@ -1007,8 +1039,7 @@ impl<'a> UpstreamPeerBuilder<'a> {
             hasher.finish()
         };
 
-        // Use a simple probabilistic approach: only log if this is likely the first time
-        if LOGGED_COUNTER.fetch_add(1, Ordering::Relaxed) % 1000 < 10 || hash % 1000 == 0 {
+        if should_log_upstream_config(hash) {
             tracing::info!(
                 upstream = %upstream_name.0,
                 idle_timeout_ms = ?peer.options.idle_timeout.map(|d| d.as_millis()),

crates/pavis/src/retry.rs

@@ -10,10 +10,53 @@
 use crate::telemetry::metrics::MetricsRegistry;
 use pavis_core::{BackoffStrategy, BodyReplayability, MethodIdempotency, RetryPolicy, RetryReason};
 use std::sync::Arc;
+use std::sync::atomic::{AtomicU64, Ordering};
 use std::time::{Duration, Instant};
 use tokio::time::sleep;
 use tracing::{debug, warn};
 
+const RETRY_BODY_GLOBAL_BUDGET_BYTES: u64 = 32 * 1024 * 1024;
+static RETRY_BODY_BUDGET_USED_BYTES: AtomicU64 = AtomicU64::new(0);
+
+#[derive(Debug)]
+pub struct BudgetGuard {
+    bytes: u64,
+}
+
+impl BudgetGuard {
+    pub fn new() -> Self {
+        Self { bytes: 0 }
+    }
+
+    pub fn try_add(&mut self, extra: u64) -> bool {
+        if extra == 0 {
+            return true;
+        }
+        loop {
+            let used = RETRY_BODY_BUDGET_USED_BYTES.load(Ordering::Relaxed);
+            let new_used = used.saturating_add(extra);
+            if new_used > RETRY_BODY_GLOBAL_BUDGET_BYTES {
+                return false;
+            }
+            if RETRY_BODY_BUDGET_USED_BYTES
+                .compare_exchange(used, new_used, Ordering::SeqCst, Ordering::Relaxed)
+                .is_ok()
+            {
+                self.bytes = self.bytes.saturating_add(extra);
+                return true;
+            }
+        }
+    }
+}
+
+impl Drop for BudgetGuard {
+    fn drop(&mut self) {
+        if self.bytes > 0 {
+            RETRY_BODY_BUDGET_USED_BYTES.fetch_sub(self.bytes, Ordering::SeqCst);
+        }
+    }
+}
+
 /// Retry context tracking state across retry attempts
 pub struct RetryContext {
     /// Global deadline for entire request (including all retries)
@@ -206,6 +249,8 @@ pub struct BufferedBody {
     pub bytes: Vec<u8>,
     /// Replayability status
     pub replayability: BodyReplayability,
+    /// Global buffer budget guard (releases on drop)
+    _budget_guard: Option<BudgetGuard>,
 }
 
 impl BufferedBody {
@@ -215,8 +260,12 @@ impl BufferedBody {
         max_buffer_size: u64,
         metrics: Option<Arc<MetricsRegistry>>,
         upstream_name: &str,
+        force_streaming: bool,
+        budget_guard: Option<BudgetGuard>,
     ) -> Self {
-        let replayability = if bytes.is_empty() {
+        let replayability = if force_streaming {
+            BodyReplayability::Streaming
+        } else if bytes.is_empty() {
             BodyReplayability::Empty
         } else if bytes.len() as u64 <= max_buffer_size {
             if let Some(metrics) = &metrics {
@@ -230,6 +279,7 @@ impl BufferedBody {
         Self {
             bytes,
             replayability,
+            _budget_guard: budget_guard,
         }
     }
 
@@ -408,36 +458,36 @@ mod tests {
 
     #[test]
     fn buffered_body_empty() {
-        let body = BufferedBody::new(vec![], 1024, None, "test");
+        let body = BufferedBody::new(vec![], 1024, None, "test", false, None);
         assert_eq!(body.replayability, BodyReplayability::Empty);
         assert!(body.can_replay());
     }
 
     #[test]
     fn buffered_body_buffered() {
-        let body = BufferedBody::new(vec![1, 2, 3], 1024, None, "test");
+        let body = BufferedBody::new(vec![1, 2, 3], 1024, None, "test", false, None);
         assert_eq!(body.replayability, BodyReplayability::Buffered);
         assert!(body.can_replay());
     }
 
     #[test]
     fn buffered_body_streaming() {
         let large_body = vec![0u8; 2048];
-        let body = BufferedBody::new(large_body, 1024, None, "test");
+        let body = BufferedBody::new(large_body, 1024, None, "test", false, None);
         assert_eq!(body.replayability, BodyReplayability::Streaming);
         assert!(!body.can_replay());
     }
 
     #[test]
     fn buffered_body_handle_non_replayable_strict() {
-        let body = BufferedBody::new(vec![0u8; 2048], 1024, None, "test");
+        let body = BufferedBody::new(vec![0u8; 2048], 1024, None, "test", false, None);
         let result = body.handle_non_replayable(true);
         assert!(result.is_err());
     }
 
     #[test]
     fn buffered_body_handle_non_replayable_lenient() {
-        let body = BufferedBody::new(vec![0u8; 2048], 1024, None, "test");
+        let body = BufferedBody::new(vec![0u8; 2048], 1024, None, "test", false, None);
         let result = body.handle_non_replayable(false);
         assert!(result.is_ok());
     }