v0.34.0_1

src/proxy/body.rs

@@ -10,7 +10,7 @@ use std::{
     future::Future,
     pin::Pin,
     sync::{
-        atomic::{AtomicBool, AtomicU16, AtomicU64, Ordering},
+        atomic::{AtomicBool, AtomicU16, AtomicU64, AtomicU8, Ordering},
         Arc, Mutex,
     },
     task::{Context, Poll},
@@ -21,13 +21,44 @@ use tokio::{
     time::{sleep_until, Sleep},
 };
 
+/// Compact source tags avoid a mutex on every completed exchange.
+#[repr(u8)]
+#[derive(Clone, Copy)]
+pub enum ExchangeSource {
+    Unresolved = 0,
+    LocalResponse = 1,
+    LocalRedirect = 2,
+    WebServer = 3,
+    ReverseProxy = 4,
+}
+impl ExchangeSource {
+    fn name(value: u8) -> &'static str {
+        match value {
+            1 => "local_response",
+            2 => "local_redirect",
+            3 => "web_server",
+            4 => "reverse_proxy",
+            _ => "unresolved",
+        }
+    }
+}
+
+// 0 = unknown, 1 = Some(false), 2 = Some(true). Preserve log schema.
+fn read_effective(value: u8) -> Option<bool> {
+    match value {
+        1 => Some(false),
+        2 => Some(true),
+        _ => None,
+    }
+}
+
 pub struct Exchange {
     pub vhost_metrics: Option<Arc<crate::metrics::VhostMetrics>>,
     pub vhost: Option<String>,
     pub route: Mutex<Option<String>>,
-    pub source: Mutex<&'static str>,
-    pub cache_effective_enabled: Mutex<Option<bool>>,
-    pub transform_effective_enabled: Mutex<Option<bool>>,
+    pub source: AtomicU8,
+    pub cache_effective_enabled: AtomicU8,
+    pub transform_effective_enabled: AtomicU8,
     pub backend: Mutex<Option<Arc<crate::balance::Lease>>>,
     pub id: u64,
     pub request_id: String,
@@ -59,6 +90,17 @@ pub struct Exchange {
     pub _permit: OwnedSemaphorePermit,
 }
 impl Exchange {
+    pub fn set_source(&self, source: ExchangeSource) {
+        self.source.store(source as u8, Ordering::Relaxed);
+    }
+    pub fn set_cache_enabled(&self, enabled: bool) {
+        self.cache_effective_enabled
+            .store(if enabled { 2 } else { 1 }, Ordering::Relaxed);
+    }
+    pub fn set_transform_enabled(&self, enabled: bool) {
+        self.transform_effective_enabled
+            .store(if enabled { 2 } else { 1 }, Ordering::Relaxed);
+    }
     pub fn websocket_phase(&self, phase: &'static str) {
         if self.websocket_bootstrap.is_some() {
             *self
@@ -145,15 +187,9 @@ impl Drop for Exchange {
             .lock()
             .unwrap_or_else(|e| e.into_inner());
         let route = self.route.lock().unwrap_or_else(|e| e.into_inner()).clone();
-        let source = *self.source.lock().unwrap_or_else(|e| e.into_inner());
-        let cache_effective_enabled = *self
-            .cache_effective_enabled
-            .lock()
-            .unwrap_or_else(|e| e.into_inner());
-        let transform_effective_enabled = *self
-            .transform_effective_enabled
-            .lock()
-            .unwrap_or_else(|e| e.into_inner());
+        let source = ExchangeSource::name(self.source.load(Ordering::Relaxed));
+        let cache_effective_enabled = read_effective(self.cache_effective_enabled.load(Ordering::Relaxed));
+        let transform_effective_enabled = read_effective(self.transform_effective_enabled.load(Ordering::Relaxed));
         let client_ip = self
             .client_ip
             .lock()
@@ -213,17 +249,17 @@ impl Drop for Exchange {
             authority: None,
             matched_by: None,
             cache_reason: *self.cache_reason.lock().unwrap_or_else(|e| e.into_inner()),
-            vhost: self.vhost.clone(),
+            vhost: self.vhost.take(),
             route,
             cache_effective_enabled,
             transform_effective_enabled,
             backend,
             id: self.id,
-            request_id: self.request_id.clone(),
+            request_id: std::mem::take(&mut self.request_id),
             time_ms: now_ms(),
-            method: self.method.clone(),
-            host: self.host.clone(),
-            path: self.path.clone(),
+            method: std::mem::take(&mut self.method),
+            host: std::mem::take(&mut self.host),
+            path: std::mem::take(&mut self.path),
             http_version: self.http_version,
             transport: self.transport,
             websocket_bootstrap: self.websocket_bootstrap,
@@ -243,7 +279,7 @@ impl Drop for Exchange {
                 .unwrap_or_else(|e| e.into_inner())
                 .clone(),
             client_ip,
-            user_agent: self.user_agent.clone(),
+            user_agent: self.user_agent.take(),
             country,
             status,
             duration_ms,
@@ -262,6 +298,9 @@ pub enum Direction {
 pub struct MeteredBody {
     inner: Body,
     exchange: Option<Arc<Exchange>>,
+    /// For a response, backend selection is final before this wrapper is made.
+    /// Keep the selected lease alive and avoid locking Exchange.backend per frame.
+    response_backend: Option<Arc<crate::balance::Lease>>,
     direction: Direction,
     timer: Option<Pin<Box<Sleep>>>,
     deadline: tokio::time::Instant,
@@ -271,9 +310,13 @@ pub struct MeteredBody {
 impl MeteredBody {
     pub fn new(inner: Body, exchange: Arc<Exchange>, direction: Direction, idle: Duration) -> Self {
         let finished = inner.is_end_stream();
+        let response_backend = matches!(direction, Direction::Response)
+            .then(|| exchange.backend.lock().unwrap_or_else(|e| e.into_inner()).clone())
+            .flatten();
         Self {
             inner,
             exchange: Some(exchange),
+            response_backend,
             direction,
             // Ready/empty bodies never need a timer allocation. Keep the
             // original deadline so a delayed first poll still times out.
@@ -286,6 +329,7 @@ impl MeteredBody {
     fn finish(&mut self) {
         self.finished = true;
         self.timer.take();
+        self.response_backend.take();
         self.exchange.take();
     }
 }
@@ -305,7 +349,14 @@ impl HttpBody for MeteredBody {
             Poll::Ready(Some(Ok(frame))) => {
                 if let (Some(data), Some(exchange)) = (frame.data_ref(), this.exchange.as_ref()) {
                     let count = data.len() as u64;
-                    exchange.backend_bytes(count, matches!(this.direction, Direction::Request));
+                    match this.direction {
+                        Direction::Request => exchange.backend_bytes(count, true),
+                        Direction::Response => {
+                            if let Some(lease) = &this.response_backend {
+                                lease.bytes(count, false);
+                            }
+                        }
+                    }
                     match this.direction {
                         Direction::Request => {
                             exchange.request_bytes.fetch_add(count, Ordering::Relaxed);