Add middleware to prioritize download traffic

This new middleware layer will reject some requests as the database pool
reaches capacity. At 20% load, the `in_flight_requests` count is added
to the log output. At 70% load, all safe requests (`GET`, `HEAD`,
`OPTIONS`, `TRACE`) are rejected immediately. This will reject many
legitimate frontend requests as well, but should catch all bot traffic
(which is unlikely to send `PUT`, `POST`, or `DELETE` requests). This
filter also helps avoid rejecting frontend requests that update the
database where we don’t always provide good feedback for errors in the
UI.

Finally, at 80% load all non-download traffic is rejected. In other words, at
least 20% of database connections are reserved for handling download
traffic. By choosing to drop other requests, there should be sufficient
database connections available to avoid queuing and timeouts on
download requests.

There is some overlap with the LogConnectionPoolStatus middleware.
These may eventually be consolidated to avoid some duplicate work and to
make smarter decisions regarding the instantaneous spare capacity of
individual pools (primary vs read-only replica).

The existing middleware does give us some insight into our current
in_flight_request counts on production. Looking through the logs, it is
rare to have more than a few requests running at the same time. This is
because download requests are completed very quickly and other API
traffic accounts for about 10 requests per second.

Author: jtgeibel

src/middleware.rs

@@ -12,6 +12,7 @@ use self::log_connection_pool_status::LogConnectionPoolStatus;
 use self::static_or_continue::StaticOrContinue;
 
 pub mod app;
+mod balance_capacity;
 mod block_traffic;
 pub mod current_user;
 mod debug;
@@ -74,6 +75,27 @@ pub fn build_middleware(app: Arc<App>, endpoints: R404) -> MiddlewareBuilder {
 
     // Note: The following `m.around()` middleware is run from bottom to top
 
+    // This is currently the final middleware to run. If a middleware layer requires a database
+    // connection, it should be run after this middleware so that the potential pool usage can be
+    // tracked here.
+    //
+    // In production we currently have 2 equally sized pools (primary and a read-only replica).
+    // Because such a large portion of production traffic is for download requests (which update
+    // download counts), we consider only the primary pool here.
+    if let Ok(capacity) = env::var("DB_POOL_SIZE") {
+        if let Ok(capacity) = capacity.parse() {
+            if capacity >= 10 {
+                println!(
+                    "Enabling BalanceCapacity middleware with {} pool capacity",
+                    capacity
+                );
+                m.around(balance_capacity::BalanceCapacity::new(capacity))
+            } else {
+                println!("BalanceCapacity middleware not enabled. DB_POOL_SIZE is too low.");
+            }
+        }
+    }
+
     // Serve the static files in the *dist* directory, which are the frontend assets.
     // Not needed for the backend tests.
     if env != Env::Test {

src/middleware/balance_capacity.rs

@@ -0,0 +1,100 @@
+//! Reject certain requests as instance load reaches capacity.
+//!
+//! The primary goal of this middleware is to avoid starving the download endpoint of resources.
+//! When bots send many parallel requests that run slow database queries, download requests may
+//! block and eventually timeout waiting for a database connection.
+//!
+//! Bots must continue to respect our crawler policy, but until we can manually block bad clients
+//! we should avoid dropping download requests even if that means rejecting some legitimate
+//! requests to other endpoints.
+
+use std::sync::atomic::{AtomicUsize, Ordering};
+
+use super::prelude::*;
+use conduit::{RequestExt, StatusCode};
+
+#[derive(Default)]
+pub(super) struct BalanceCapacity {
+    handler: Option<Box<dyn Handler>>,
+    capacity: usize,
+    in_flight_requests: AtomicUsize,
+}
+
+impl BalanceCapacity {
+    pub fn new(capacity: usize) -> Self {
+        Self {
+            handler: None,
+            capacity,
+            in_flight_requests: AtomicUsize::new(0),
+        }
+    }
+}
+
+impl AroundMiddleware for BalanceCapacity {
+    fn with_handler(&mut self, handler: Box<dyn Handler>) {
+        self.handler = Some(handler);
+    }
+}
+
+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 at 20% capacity
+        if load >= 20 {
+            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);
+        }
+
+        // Reject read-only requests after reaching 70% load. Bots are likely to send only safe
+        // requests and this helps prioritize requests that users may be reluctant to retry.
+        if load >= 70 && request.method().is_safe() {
+            return over_capcity_response();
+        }
+
+        // At 80% load, all non-download requests are rejected
+        if load >= 80 {
+            return over_capcity_response();
+        }
+
+        handler.call(request)
+    }
+}
+
+fn over_capcity_response() -> AfterResult {
+    // TODO: Generate an alert so we can investigate
+    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)
+}
+
+// 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,
+}
+
+impl<'a> RequestCounter<'a> {
+    fn add_one(counter: &'a AtomicUsize) -> (Self, usize) {
+        let previous = counter.fetch_add(1, Ordering::SeqCst);
+        (Self { counter }, previous + 1)
+    }
+}
+
+impl<'a> Drop for RequestCounter<'a> {
+    fn drop(&mut self) {
+        self.counter.fetch_sub(1, Ordering::SeqCst);
+    }
+}