Update report-only mode to include (proposed) rejection reason

Author: jtgeibel

src/middleware/balance_capacity.rs

@@ -37,6 +37,32 @@ impl BalanceCapacity {
             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 {
@@ -49,50 +75,33 @@ impl Handler for BalanceCapacity {
     fn call(&self, request: &mut dyn RequestExt) -> AfterResult {
         // 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 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);
         }
 
-        // In report-only mode we serve all requests and only enforce the logging limit above
-        if self.report_only {
-            return handler.call(request);
-        }
-
         // Download requests are always accepted
         if request.path().starts_with("/api/v1/crates/") && request.path().ends_with("/download") {
-            return handler.call(request);
+            return self.handle(request);
         }
 
         // 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) = env::var(name) {
         value.parse().unwrap_or(default)