Auto merge of #2495 - jtgeibel:web-capacity-report-only, r=pietroalbini

Add a report-only mode to the `BalanceCapacity` middleware

This adds a single environment variable that can be quickly set or
removed to toggle capacity enforcement. If this variable is set to any
value (including the empty string) then no requests will be rejected due
to load.

This will allow us to evolve the heuristics used in this middleware to
find the right balance before turning on enforcement.

This also increases the default logging percentage to a higher default.

r? @pietroalbini

Author: bors

src/middleware/balance_capacity.rs

@@ -8,6 +8,7 @@
 //! we should avoid dropping download requests even if that means rejecting some legitimate
 //! requests to other endpoints.
 
+use std::env;
 use std::sync::atomic::{AtomicUsize, Ordering};
 
 use super::prelude::*;
@@ -17,7 +18,8 @@ use conduit::{RequestExt, StatusCode};
 pub(super) struct BalanceCapacity {
     handler: Option<Box<dyn Handler>>,
     capacity: usize,
-    in_flight_requests: AtomicUsize,
+    in_flight_non_dl_requests: AtomicUsize,
+    report_only: bool,
     log_at_percentage: usize,
     throttle_at_percentage: usize,
     dl_only_at_percentage: usize,
@@ -28,12 +30,39 @@ impl BalanceCapacity {
         Self {
             handler: None,
             capacity,
-            in_flight_requests: AtomicUsize::new(0),
-            log_at_percentage: read_env_percentage("WEB_CAPACITY_LOG_PCT", 20),
+            in_flight_non_dl_requests: AtomicUsize::new(0),
+            report_only: env::var("WEB_CAPACITY_REPORT_ONLY").ok().is_some(),
+            log_at_percentage: read_env_percentage("WEB_CAPACITY_LOG_PCT", 50),
             throttle_at_percentage: read_env_percentage("WEB_CAPACITY_THROTTLE_PCT", 70),
             dl_only_at_percentage: read_env_percentage("WEB_CAPACITY_DL_ONLY_PCT", 80),
         }
     }
+
+    /// Handle a request normally
+    fn handle(&self, request: &mut dyn RequestExt) -> AfterResult {
+        self.handler.as_ref().unwrap().call(request)
+    }
+
+    /// Handle a request when load exceeds a threshold
+    ///
+    /// In report-only mode, log metadata is added but the request is still served. Otherwise,
+    /// the request is rejected with a service unavailable response.
+    fn handle_high_load(&self, request: &mut dyn RequestExt, note: &str) -> AfterResult {
+        if self.report_only {
+            // In report-only mode we serve all requests but add log metadata
+            super::log_request::add_custom_metadata(request, "would_reject", note);
+            self.handle(request)
+        } else {
+            // Reject the request
+            super::log_request::add_custom_metadata(request, "cause", note);
+            let body = "Service temporarily unavailable";
+            Response::builder()
+                .status(StatusCode::SERVICE_UNAVAILABLE)
+                .header(header::CONTENT_LENGTH, body.len())
+                .body(Body::from_static(body.as_bytes()))
+                .map_err(box_error)
+        }
+    }
 }
 
 impl AroundMiddleware for BalanceCapacity {
@@ -44,58 +73,43 @@ impl AroundMiddleware for BalanceCapacity {
 
 impl Handler for BalanceCapacity {
     fn call(&self, request: &mut dyn RequestExt) -> AfterResult {
+        // Download requests are always accepted and do not affect the request count
+        if request.path().starts_with("/api/v1/crates/") && request.path().ends_with("/download") {
+            return self.handle(request);
+        }
+
         // The _drop_on_exit ensures the counter is decremented for all exit paths (including panics)
-        let (_drop_on_exit, count) = RequestCounter::add_one(&self.in_flight_requests);
-        let handler = self.handler.as_ref().unwrap();
+        let (_drop_on_exit, count) = RequestCounter::add_one(&self.in_flight_non_dl_requests);
         let load = 100 * count / self.capacity;
 
         // Begin logging request count so early stages of load increase can be located
         if load >= self.log_at_percentage {
-            super::log_request::add_custom_metadata(request, "in_flight_requests", count);
-        }
-
-        // Download requests are always accepted
-        if request.path().starts_with("/api/v1/crates/") && request.path().ends_with("/download") {
-            return handler.call(request);
+            super::log_request::add_custom_metadata(request, "in_flight_non_dl_requests", count);
         }
 
         // Reject read-only requests as load nears capacity. Bots are likely to send only safe
         // requests and this helps prioritize requests that users may be reluctant to retry.
         if load >= self.throttle_at_percentage && request.method().is_safe() {
-            return over_capacity_response(request);
+            return self.handle_high_load(request, "over capacity (throttle)");
         }
 
         // As load reaches capacity, all non-download requests are rejected
         if load >= self.dl_only_at_percentage {
-            return over_capacity_response(request);
+            return self.handle_high_load(request, "over capacity (download only)");
         }
 
-        handler.call(request)
+        self.handle(request)
     }
 }
 
-fn over_capacity_response(request: &mut dyn RequestExt) -> AfterResult {
-    // TODO: Generate an alert so we can investigate
-    super::log_request::add_custom_metadata(request, "cause", "over capacity");
-    let body = "Service temporarily unavailable";
-    Response::builder()
-        .status(StatusCode::SERVICE_UNAVAILABLE)
-        .header(header::CONTENT_LENGTH, body.len())
-        .body(Body::from_static(body.as_bytes()))
-        .map_err(box_error)
-}
-
 fn read_env_percentage(name: &str, default: usize) -> usize {
-    if let Ok(value) = std::env::var(name) {
+    if let Ok(value) = env::var(name) {
         value.parse().unwrap_or(default)
     } else {
         default
     }
 }
 
-// FIXME(JTG): I've copied the following from my `conduit-hyper` crate.  Once we transition from
-// `civet`, we could pass the in_flight_request count from `condut-hyper` via a request extension.
-
 /// A struct that stores a reference to an atomic counter so it can be decremented when dropped
 struct RequestCounter<'a> {
     counter: &'a AtomicUsize,