0.34.6_repo_cleanup

src/proxy/body/mod.rs

@@ -0,0 +1,511 @@
+use crate::{
+    logging::{now_ms, LogEntry, Logs},
+    metrics::Metrics,
+};
+use axum::body::Body;
+use bytes::Bytes;
+use http_body::{Body as HttpBody, Frame, SizeHint};
+use std::{
+    future::Future,
+    pin::Pin,
+    sync::{
+        atomic::{AtomicBool, AtomicU16, AtomicU64, AtomicU8, Ordering},
+        Arc, Mutex,
+    },
+    task::{Context, Poll},
+    time::{Duration, Instant},
+};
+use tokio::{
+    sync::OwnedSemaphorePermit,
+    time::{sleep_until, Sleep},
+};
+
+#[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",
+        }
+    }
+}
+
+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: 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,
+    pub method: String,
+    pub host: String,
+    pub path: String,
+    pub http_version: &'static str,
+    pub transport: &'static str,
+    pub websocket_bootstrap: Option<&'static str>,
+    pub websocket_phase: Mutex<Option<&'static str>>,
+    pub websocket_error: Mutex<Option<&'static str>>,
+    pub h3_reset_code: Mutex<Option<&'static str>>,
+    pub upstream_status: AtomicU16,
+    pub upstream_error: Mutex<Option<String>>,
+    pub client_ip: Mutex<Option<String>>,
+    pub user_agent: Option<String>,
+    pub country: Mutex<Option<String>>,
+    pub started: Instant,
+    pub status: AtomicU16,
+    pub cache_status: AtomicU16,
+    pub cache_reason: Mutex<&'static str>,
+    pub failed: AtomicBool,
+    pub aborted: AtomicBool,
+    pub request_bytes: AtomicU64,
+    pub response_bytes: AtomicU64,
+    pub metrics: Arc<Metrics>,
+    pub logs: Arc<Logs>,
+    pub traffic: Arc<crate::traffic::Analytics>,
+    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
+                .websocket_phase
+                .lock()
+                .unwrap_or_else(|e| e.into_inner()) = Some(phase);
+        }
+    }
+    pub fn websocket_error(&self, phase: &'static str, error: &'static str) {
+        if self.websocket_bootstrap.is_some() {
+            *self
+                .websocket_phase
+                .lock()
+                .unwrap_or_else(|e| e.into_inner()) = Some(phase);
+            *self
+                .websocket_error
+                .lock()
+                .unwrap_or_else(|e| e.into_inner()) = Some(error);
+        }
+    }
+    pub fn h3_reset(&self, code: &'static str) {
+        *self.h3_reset_code.lock().unwrap_or_else(|e| e.into_inner()) = Some(code);
+    }
+    pub fn backend_bytes(&self, bytes: u64, incoming: bool) {
+        if let Some(backend) = self
+            .backend
+            .lock()
+            .unwrap_or_else(|e| e.into_inner())
+            .as_ref()
+        {
+            backend.bytes(bytes, incoming);
+        }
+    }
+    fn backend_error(&self, reason: &'static str) {
+        if let Some(backend) = self
+            .backend
+            .lock()
+            .unwrap_or_else(|e| e.into_inner())
+            .as_ref()
+        {
+            backend.error(reason, self.started.elapsed());
+        }
+    }
+}
+impl Drop for Exchange {
+    fn drop(&mut self) {
+        let status = self.status.load(Ordering::Relaxed);
+        let failed = self.failed.load(Ordering::Relaxed);
+        let aborted = self.aborted.load(Ordering::Relaxed);
+        let duration = self.started.elapsed();
+        if let Some(m) = &self.vhost_metrics {
+            m.complete(
+                status,
+                failed || aborted,
+                duration,
+                self.request_bytes.load(Ordering::Relaxed),
+                self.response_bytes.load(Ordering::Relaxed),
+            );
+        }
+        self.metrics.active.fetch_sub(1, Ordering::Relaxed);
+        self.metrics.completed.fetch_add(1, Ordering::Relaxed);
+        self.metrics
+            .duration_us
+            .fetch_add(duration.as_micros() as u64, Ordering::Relaxed);
+        if failed || aborted || status >= 500 {
+            self.metrics.errors.fetch_add(1, Ordering::Relaxed);
+        }
+        match status / 100 {
+            1 => &self.metrics.status_1xx,
+            2 => &self.metrics.status_2xx,
+            3 => &self.metrics.status_3xx,
+            4 => &self.metrics.status_4xx,
+            _ => &self.metrics.status_5xx,
+        }
+        .fetch_add(1, Ordering::Relaxed);
+        let backend = self
+            .backend
+            .get_mut()
+            .unwrap_or_else(|e| e.into_inner())
+            .as_ref()
+            .map(|l| l.member.backend.id.clone());
+        let websocket_phase = *self
+            .websocket_phase
+            .get_mut()
+            .unwrap_or_else(|e| e.into_inner());
+        let route = self
+            .route
+            .get_mut()
+            .unwrap_or_else(|e| e.into_inner())
+            .take();
+        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
+            .get_mut()
+            .unwrap_or_else(|e| e.into_inner())
+            .take();
+        let country = self
+            .country
+            .get_mut()
+            .unwrap_or_else(|e| e.into_inner())
+            .take();
+        let request_bytes = self.request_bytes.load(Ordering::Relaxed);
+        let response_bytes = self.response_bytes.load(Ordering::Relaxed);
+        let duration_ms = duration.as_millis() as u64;
+        let cache = match self.cache_status.load(Ordering::Relaxed) {
+            1 => "HIT",
+            2 => "MISS",
+            3 => "BYPASS",
+            _ => "-",
+        };
+        let outcome = if failed {
+            "body_error"
+        } else if aborted {
+            "cancelled"
+        } else if status >= 500 {
+            "proxy_or_upstream_error"
+        } else if status == 101 || websocket_phase == Some("tunnel_closed") {
+            "websocket_closed"
+        } else {
+            "complete"
+        };
+        if let Some(vhost) = self.vhost.as_ref() {
+            self.traffic.push_lazy(|| crate::traffic::Sample {
+                vhost: vhost.clone(),
+                route: route.clone(),
+                method: self.method.clone(),
+                path: self.path.clone(),
+                client_ip: client_ip.clone(),
+                status,
+                duration_ms,
+                request_bytes,
+                response_bytes,
+                cache,
+                source,
+                user_agent: self.user_agent.clone(),
+                country: country.clone(),
+            });
+        }
+        self.logs.push(LogEntry {
+            event: "request",
+            protocol_error: None,
+            reason: None,
+            error_detail: None,
+            error: None,
+            local_addr: None,
+            tls_sni: None,
+            sni: None,
+            authority: None,
+            matched_by: None,
+            cache_reason: *self
+                .cache_reason
+                .get_mut()
+                .unwrap_or_else(|e| e.into_inner()),
+            vhost: self.vhost.take(),
+            route,
+            cache_effective_enabled,
+            transform_effective_enabled,
+            backend,
+            id: self.id,
+            request_id: std::mem::take(&mut self.request_id),
+            time_ms: now_ms(),
+            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,
+            websocket_phase,
+            websocket_error: *self
+                .websocket_error
+                .get_mut()
+                .unwrap_or_else(|e| e.into_inner()),
+            h3_reset_code: *self
+                .h3_reset_code
+                .get_mut()
+                .unwrap_or_else(|e| e.into_inner()),
+            upstream_status: match self.upstream_status.load(Ordering::Relaxed) {
+                0 => None,
+                value => Some(value),
+            },
+            upstream_error: self
+                .upstream_error
+                .get_mut()
+                .unwrap_or_else(|e| e.into_inner())
+                .take(),
+            client_ip,
+            user_agent: self.user_agent.take(),
+            country,
+            status,
+            duration_ms,
+            cache,
+            request_bytes,
+            response_bytes,
+            outcome,
+        });
+    }
+}
+
+#[derive(Clone, Copy)]
+pub enum Direction {
+    Request,
+    Response,
+}
+pub struct MeteredBody {
+    inner: Body,
+    exchange: Option<Arc<Exchange>>,
+    response_backend: Option<Arc<crate::balance::Lease>>,
+    request_backend: Option<Arc<crate::balance::Lease>>,
+    direction: Direction,
+    timer: Option<Pin<Box<Sleep>>>,
+    deadline: tokio::time::Instant,
+    idle: Duration,
+    finished: bool,
+}
+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,
+            request_backend: None,
+            direction,
+            timer: None,
+            deadline: tokio::time::Instant::now() + idle,
+            idle,
+            finished,
+        }
+    }
+    fn finish(&mut self) {
+        self.finished = true;
+        self.timer.take();
+        self.response_backend.take();
+        self.request_backend.take();
+        self.exchange.take();
+    }
+}
+impl HttpBody for MeteredBody {
+    type Data = Bytes;
+    type Error = axum::Error;
+    fn poll_frame(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<Result<Frame<Bytes>, Self::Error>>> {
+        let this = self.get_mut();
+        if this.finished {
+            this.exchange.take();
+            return Poll::Ready(None);
+        }
+        match Pin::new(&mut this.inner).poll_frame(cx) {
+            Poll::Ready(Some(Ok(frame))) => {
+                if let (Some(data), Some(exchange)) = (frame.data_ref(), this.exchange.as_ref()) {
+                    let count = data.len() as u64;
+                    match this.direction {
+                        Direction::Request => {
+                            if this.request_backend.is_none() {
+                                this.request_backend = exchange
+                                    .backend
+                                    .lock()
+                                    .unwrap_or_else(|e| e.into_inner())
+                                    .clone();
+                            }
+                            if let Some(lease) = &this.request_backend {
+                                lease.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);
+                            exchange
+                                .metrics
+                                .bytes_in
+                                .fetch_add(count, Ordering::Relaxed);
+                        }
+                        Direction::Response => {
+                            exchange.response_bytes.fetch_add(count, Ordering::Relaxed);
+                            exchange
+                                .metrics
+                                .bytes_out
+                                .fetch_add(count, Ordering::Relaxed);
+                        }
+                    }
+                }
+                if this.inner.is_end_stream() {
+                    this.finish();
+                } else {
+                    this.deadline = tokio::time::Instant::now() + this.idle;
+                    if let Some(timer) = this.timer.as_mut() {
+                        timer.as_mut().reset(this.deadline);
+                    }
+                }
+                Poll::Ready(Some(Ok(frame)))
+            }
+            Poll::Ready(Some(Err(error))) => {
+                if let Some(exchange) = &this.exchange {
+                    exchange.failed.store(true, Ordering::Relaxed);
+                    exchange.backend_error(match this.direction {
+                        Direction::Request => "send_error",
+                        Direction::Response => "read_error",
+                    });
+                }
+                this.finish();
+                Poll::Ready(Some(Err(error)))
+            }
+            Poll::Ready(None) => {
+                this.finish();
+                Poll::Ready(None)
+            }
+            Poll::Pending => {
+                let timer = this
+                    .timer
+                    .get_or_insert_with(|| Box::pin(sleep_until(this.deadline)));
+                if timer.as_mut().poll(cx).is_ready() {
+                    if let Some(exchange) = &this.exchange {
+                        exchange.failed.store(true, Ordering::Relaxed);
+                        exchange.backend_error(match this.direction {
+                            Direction::Request => "send_timeout",
+                            Direction::Response => "read_timeout",
+                        });
+                    }
+                    this.finish();
+                    Poll::Ready(Some(Err(axum::Error::new(std::io::Error::new(
+                        std::io::ErrorKind::TimedOut,
+                        "body idle timeout",
+                    )))))
+                } else {
+                    Poll::Pending
+                }
+            }
+        }
+    }
+    fn is_end_stream(&self) -> bool {
+        self.finished || self.inner.is_end_stream()
+    }
+    fn size_hint(&self) -> SizeHint {
+        self.inner.size_hint()
+    }
+}
+impl Drop for MeteredBody {
+    fn drop(&mut self) {
+        if !self.finished
+            && !self.inner.is_end_stream()
+            && matches!(self.direction, Direction::Response)
+        {
+            if let Some(exchange) = &self.exchange {
+                exchange.aborted.store(true, Ordering::Relaxed);
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+include!("tests.rs");
+
+pub struct ReplayBody {
+    head: Option<http_body::Frame<bytes::Bytes>>,
+    tail: axum::body::Body,
+}
+impl ReplayBody {
+    pub fn new(head: http_body::Frame<bytes::Bytes>, tail: axum::body::Body) -> Self {
+        Self {
+            head: Some(head),
+            tail,
+        }
+    }
+}
+impl http_body::Body for ReplayBody {
+    type Data = bytes::Bytes;
+    type Error = axum::Error;
+    fn poll_frame(
+        mut self: std::pin::Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<Option<Result<http_body::Frame<Self::Data>, Self::Error>>> {
+        if let Some(head) = self.head.take() {
+            return std::task::Poll::Ready(Some(Ok(head)));
+        }
+        std::pin::Pin::new(&mut self.tail).poll_frame(cx)
+    }
+    fn is_end_stream(&self) -> bool {
+        self.head.is_none() && self.tail.is_end_stream()
+    }
+    fn size_hint(&self) -> http_body::SizeHint {
+        let mut hint = http_body::SizeHint::new();
+        if self.head.is_none() {
+            return self.tail.size_hint();
+        }
+        if let Some(data) = self.head.as_ref().and_then(|f| f.data_ref()) {
+            hint.set_lower(data.len() as u64);
+        }
+        hint
+    }
+}

src/proxy/body/tests.rs

@@ -168,8 +168,6 @@ mod tests {
             Direction::Request,
             Duration::from_secs(1),
         );
-        // Move the saved initial deadline into the past without wall-clock
-        // sleeps: the first Pending must create an already-expired timer.
         body.deadline = tokio::time::Instant::now() - Duration::from_secs(1);
         assert!(body.timer.is_none());
         let result = tokio::time::timeout(Duration::from_millis(200), body.frame()).await;