Refactor pool key sampling logic: improve readability of environment variable parsing and formatting

fabian4 · fabian4 · commit 135597bcbb28 · 2026-02-04T14:11:02.000+08:00
diff --git a/crates/pavis/src/proxy/service/io.rs b/crates/pavis/src/proxy/service/io.rs
@@ -33,6 +33,9 @@ static SNI_FRAGMENT_WARN_COUNTER: std::sync::atomic::AtomicU64 =
 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();
+static POOL_KEY_SAMPLE_COUNTER: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
+static POOL_KEY_SAMPLE_RATE: OnceLock<u64> = OnceLock::new();
+const DEFAULT_POOL_KEY_SAMPLE_RATE: u64 = 16;
 
 fn should_log_upstream_config(hash: u64) -> bool {
     if !LOGGED_UPSTREAM_CHECK_COUNTER
@@ -48,6 +51,23 @@ fn should_log_upstream_config(hash: u64) -> bool {
     false
 }
 
+fn should_sample_pool_key() -> bool {
+    let rate =
+        *POOL_KEY_SAMPLE_RATE.get_or_init(|| match std::env::var("PAVIS_POOL_KEY_SAMPLE_RATE") {
+            Ok(value) => value
+                .parse::<u64>()
+                .ok()
+                .filter(|v| *v > 0)
+                .unwrap_or(DEFAULT_POOL_KEY_SAMPLE_RATE),
+            Err(_) => DEFAULT_POOL_KEY_SAMPLE_RATE,
+        });
+    if rate == 1 {
+        return true;
+    }
+    let count = POOL_KEY_SAMPLE_COUNTER.fetch_add(1, Ordering::Relaxed);
+    count % rate == 0
+}
+
 #[async_trait]
 impl ProxyHttp for Proxy {
     type CTX = RouterContext;
@@ -875,7 +895,9 @@ impl<'a> UpstreamPeerBuilder<'a> {
         }
 
         let sni_label = sni_value.as_ref().map(|name| name.0.as_str()).unwrap_or("");
-        if let Some(tracker) = self.telemetry.pool_key_tracker.as_ref() {
+        if let Some(tracker) = self.telemetry.pool_key_tracker.as_ref()
+            && should_sample_pool_key()
+        {
             let snapshot = tracker.record(
                 upstream_name.0.as_str(),
                 reuse_key_hash(&addr, sni_label, verify_mode, cert),
diff --git a/crates/pavis/src/router/matcher.rs b/crates/pavis/src/router/matcher.rs
@@ -118,17 +118,8 @@ pub(crate) fn match_request<'a>(
                         }
 
                         if !matches_method(&route.matcher.method, method) {
-                            println!(
-                                "DEBUG: Method mismatch for route {}: req={} predicate={:?}",
-                                compiled.index, method, route.matcher.method
-                            );
                             stats.record_method_miss();
                             continue;
-                        } else {
-                            println!(
-                                "DEBUG: Method MATCH for route {}: req={} predicate={:?}",
-                                compiled.index, method, route.matcher.method
-                            );
                         }
 
                         if !matches_headers(
diff --git a/crates/pavis/src/telemetry/access_log.rs b/crates/pavis/src/telemetry/access_log.rs
@@ -13,7 +13,7 @@ use tokio::sync::mpsc;
 use tokio::time::Duration;
 
 pub struct AccessLog {
-    tx: mpsc::Sender<LogEntry>,
+    tx: Option<mpsc::Sender<LogEntry>>,
     enabled: bool,
     metrics: Mutex<Option<Arc<MetricsRegistry>>>,
 }
@@ -32,6 +32,9 @@ impl Service for AccessLogWorker {
         mut shutdown: tokio::sync::watch::Receiver<bool>,
         _threads: usize,
     ) {
+        if matches!(self.config, AccessLogPolicy::Disabled) || self.rx.is_none() {
+            return;
+        }
         let mut rx = self.rx.take().expect("Worker started twice");
 
         let mut file_writer = if let AccessLogPolicy::File(path) = &self.config {
@@ -117,12 +120,17 @@ impl Service for AccessLogWorker {
 
 impl AccessLog {
     pub fn new(config: &AccessLogPolicy) -> (Self, AccessLogWorker) {
-        let (tx, rx) = mpsc::channel::<LogEntry>(access_log_channel_capacity());
         let enabled = *config != AccessLogPolicy::Disabled;
+        let (tx, rx) = if enabled {
+            let (tx, rx) = mpsc::channel::<LogEntry>(access_log_channel_capacity());
+            (Some(tx), Some(rx))
+        } else {
+            (None, None)
+        };
         let throttle_ms = access_log_throttle_ms();
 
         let worker = AccessLogWorker {
-            rx: Some(rx),
+            rx,
             config: config.clone(),
             throttle_ms,
         };
@@ -159,6 +167,9 @@ impl AccessLog {
         if !enabled {
             return;
         }
+        let Some(tx) = &self.tx else {
+            return;
+        };
 
         let req = session.req_header();
         let method = req.method.clone();
@@ -212,7 +223,7 @@ impl AccessLog {
         };
 
         // Non-blocking send (lossy if full)
-        if let Err(mpsc::error::TrySendError::Full(_)) = self.tx.try_send(entry)
+        if let Err(mpsc::error::TrySendError::Full(_)) = tx.try_send(entry)
             && let Some(handle) = self
                 .metrics
                 .lock()
@@ -367,7 +378,9 @@ mod tests {
             upstream_duration_ms: Some(50),
         };
         let expected = format_log_line(&entry);
-        let _ = access_log.tx.try_send(entry);
+        if let Some(tx) = &access_log.tx {
+            let _ = tx.try_send(entry);
+        }
 
         // Run worker briefly
         let (shutdown_tx, shutdown_rx) = watch::channel(false);
diff --git a/crates/pavis/src/telemetry/metrics.rs b/crates/pavis/src/telemetry/metrics.rs
@@ -1,6 +1,6 @@
 use crate::router::MatchVerdict;
 use async_trait::async_trait;
-use metrics::{counter, gauge, histogram};
+use metrics::{SharedString, counter, gauge, histogram};
 use metrics_exporter_prometheus::PrometheusBuilder;
 use pingora::services::Service;
 use std::collections::{HashMap, VecDeque};
@@ -48,6 +48,7 @@ impl<T: MetricsTransport> PrometheusEndpoint<T> {
             Ok(handle) => {
                 let registry = MetricsRegistry {
                     _handle: Arc::new(handle),
+                    labels: Arc::new(MetricLabels::new()),
                 };
                 (
                     Self {
@@ -196,9 +197,12 @@ async fn serve_metrics(
 #[derive(Clone)]
 pub struct MetricsRegistry {
     _handle: Arc<metrics_exporter_prometheus::PrometheusHandle>,
+    labels: Arc<MetricLabels>,
 }
 
 pub const POOL_KEY_CARDINALITY_CAP: usize = 1024;
+const METRIC_LABEL_CACHE_CAP: usize = 1024;
+const METRIC_ROUTE_LABEL_CACHE_CAP: usize = 4096;
 const METRICS_REQUEST_LINE_LIMIT_BYTES: usize = 4096;
 const METRICS_HEADER_LIMIT_BYTES: usize = 16 * 1024;
 const METRICS_READ_TIMEOUT: Duration = Duration::from_secs(5);
@@ -215,6 +219,87 @@ struct BoundedKeySet {
     order: VecDeque<(u64, Instant)>,
 }
 
+struct MetricLabels {
+    common: LabelCache,
+    route: LabelCache,
+    status: Mutex<HashMap<u16, SharedString>>,
+}
+
+impl MetricLabels {
+    fn new() -> Self {
+        Self {
+            common: LabelCache::new(METRIC_LABEL_CACHE_CAP),
+            route: LabelCache::new(METRIC_ROUTE_LABEL_CACHE_CAP),
+            status: Mutex::new(HashMap::new()),
+        }
+    }
+
+    fn common(&self, value: &str) -> SharedString {
+        self.common.get(value)
+    }
+
+    fn route(&self, value: &str) -> SharedString {
+        self.route.get(value)
+    }
+
+    fn status(&self, value: u16) -> SharedString {
+        let mut guard = self
+            .status
+            .lock()
+            .expect("metrics status cache lock poisoned");
+        if let Some(existing) = guard.get(&value) {
+            return existing.clone();
+        }
+        let shared = SharedString::from(value.to_string());
+        guard.insert(value, shared.clone());
+        shared
+    }
+}
+
+struct LabelCache {
+    cap: usize,
+    inner: Mutex<LabelCacheInner>,
+}
+
+struct LabelCacheInner {
+    map: HashMap<String, SharedString>,
+    order: VecDeque<String>,
+}
+
+impl LabelCache {
+    fn new(cap: usize) -> Self {
+        Self {
+            cap,
+            inner: Mutex::new(LabelCacheInner {
+                map: HashMap::new(),
+                order: VecDeque::new(),
+            }),
+        }
+    }
+
+    fn get(&self, value: &str) -> SharedString {
+        let mut guard = self
+            .inner
+            .lock()
+            .expect("metrics label cache lock poisoned");
+        if let Some(existing) = guard.map.get(value) {
+            return existing.clone();
+        }
+        let shared = SharedString::from(value.to_string());
+        let key = value.to_string();
+        guard.map.insert(key.clone(), shared.clone());
+        guard.order.push_back(key);
+        while guard.map.len() > self.cap {
+            if let Some(evicted) = guard.order.pop_front() {
+                guard.map.remove(&evicted);
+            } else {
+                break;
+            }
+        }
+        shared
+    }
+}
+
 impl BoundedKeySet {
     fn new(cap: usize) -> Self {
         Self {
@@ -369,36 +454,43 @@ impl MetricsRegistry {
         upstream: &str,
         duration_secs: f64,
     ) {
+        let method = self.labels.common(method);
+        let route = self.labels.route(route_pattern);
+        let status = self.labels.status(status);
+        let upstream = self.labels.common(upstream);
+
         counter!(
             "pavis_http_requests_total",
-            "method" => method.to_string(),
-            "route" => route_pattern.to_string(),
-            "status" => status.to_string(),
-            "upstream" => upstream.to_string(),
+            "method" => method,
+            "route" => route,
+            "status" => status.clone(),
+            "upstream" => upstream.clone(),
         )
         .increment(1);
 
         histogram!(
             "pavis_http_request_duration_seconds",
-            "method" => method.to_string(),
-            "route" => route_pattern.to_string(),
-            "status" => status.to_string(),
-            "upstream" => upstream.to_string(),
+            "method" => method,
+            "route" => route,
+            "status" => status,
+            "upstream" => upstream,
         )
         .record(duration_secs);
     }
 
     pub fn record_upstream_request(&self, upstream: &str, status: u16, duration_secs: f64) {
+        let upstream = self.labels.common(upstream);
+        let status = self.labels.status(status);
         counter!(
             "pavis_upstream_requests_total",
-            "upstream" => upstream.to_string(),
-            "status" => status.to_string(),
+            "upstream" => upstream.clone(),
+            "status" => status.clone(),
         )
         .increment(1);
 
         histogram!(
             "pavis_upstream_request_duration_seconds",
-            "upstream" => upstream.to_string(),
+            "upstream" => upstream,
         )
         .record(duration_secs);
     }
@@ -413,41 +505,47 @@ impl MetricsRegistry {
     }
 
     pub fn record_pool_size(&self, upstream: &str, size: f64) {
-        gauge!("pavis_upstream_pool_size", "upstream" => upstream.to_string()).set(size);
+        let upstream = self.labels.common(upstream);
+        gauge!("pavis_upstream_pool_size", "upstream" => upstream).set(size);
     }
 
     pub fn record_pool_key_cardinality(&self, upstream: &str, cardinality: usize, saturated: bool) {
+        let upstream = self.labels.common(upstream);
         let reported = if saturated {
             (POOL_KEY_CARDINALITY_CAP + 1) as f64
         } else {
             cardinality as f64
         };
         gauge!(
             "pavis_upstream_pool_key_cardinality_approx",
-            "upstream" => upstream.to_string()
+            "upstream" => upstream
         )
         .set(reported);
     }
 
     pub fn record_connection_reused(&self, upstream: &str) {
+        let upstream = self.labels.common(upstream);
         counter!(
             "pavis_upstream_connection_reused_total",
-            "upstream" => upstream.to_string()
+            "upstream" => upstream
         )
         .increment(1);
     }
 
     pub fn record_connection_new(&self, upstream: &str, reason: &str) {
+        let upstream = self.labels.common(upstream);
+        let reason = self.labels.common(reason);
         counter!(
             "pavis_upstream_connection_new_total",
-            "upstream" => upstream.to_string(),
-            "reason" => reason.to_string()
+            "upstream" => upstream,
+            "reason" => reason
         )
         .increment(1);
     }
 
     pub fn update_config_stats(&self, version: &str, size_bytes: u64) {
-        gauge!("pavis_runtime_config_version", "version" => version.to_string()).set(1.0);
+        let version = self.labels.common(version);
+        gauge!("pavis_runtime_config_version", "version" => version).set(1.0);
         gauge!("pavis_runtime_config_size_bytes").set(size_bytes as f64);
         gauge!("pavis_runtime_reload_last_timestamp").set(
             std::time::SystemTime::now()
@@ -462,18 +560,21 @@ impl MetricsRegistry {
     }
 
     pub fn record_config_validation(&self, result: &str, reason: &str) {
+        let result = self.labels.common(result);
+        let reason = self.labels.common(reason);
         counter!(
             "pavis_config_validation_total",
-            "result" => result.to_string(),
-            "reason" => reason.to_string(),
+            "result" => result,
+            "reason" => reason,
         )
         .increment(1);
     }
 
     pub fn record_config_apply(&self, result: &str) {
+        let result = self.labels.common(result);
         counter!(
             "pavis_config_apply_total",
-            "result" => result.to_string(),
+            "result" => result,
         )
         .increment(1);
     }
diff --git a/crates/pavis/src/upstream/cluster/pool.rs b/crates/pavis/src/upstream/cluster/pool.rs
