proxy: cull http connections (#7632)

## Problem

Some HTTP client connections can stay open for quite a long time.

## Summary of changes

When there are too many HTTP client connections, pick a random
connection and gracefully cancel it.

proxy/src/serverless/cancel_set.rs

@@ -0,0 +1,102 @@
+//! A set for cancelling random http connections
+
+use std::{
+    hash::{BuildHasher, BuildHasherDefault},
+    num::NonZeroUsize,
+    time::Duration,
+};
+
+use indexmap::IndexMap;
+use parking_lot::Mutex;
+use rand::{thread_rng, Rng};
+use rustc_hash::FxHasher;
+use tokio::time::Instant;
+use tokio_util::sync::CancellationToken;
+use uuid::Uuid;
+
+type Hasher = BuildHasherDefault<FxHasher>;
+
+pub struct CancelSet {
+    shards: Box<[Mutex<CancelShard>]>,
+    // keyed by random uuid, fxhasher is fine
+    hasher: Hasher,
+}
+
+pub struct CancelShard {
+    tokens: IndexMap<uuid::Uuid, (Instant, CancellationToken), Hasher>,
+}
+
+impl CancelSet {
+    pub fn new(shards: usize) -> Self {
+        CancelSet {
+            shards: (0..shards)
+                .map(|_| {
+                    Mutex::new(CancelShard {
+                        tokens: IndexMap::with_hasher(Hasher::default()),
+                    })
+                })
+                .collect(),
+            hasher: Hasher::default(),
+        }
+    }
+
+    pub fn take(&self) -> Option<CancellationToken> {
+        for _ in 0..4 {
+            if let Some(token) = self.take_raw(thread_rng().gen()) {
+                return Some(token);
+            }
+            tracing::trace!("failed to get cancel token");
+        }
+        None
+    }
+
+    pub fn take_raw(&self, rng: usize) -> Option<CancellationToken> {
+        NonZeroUsize::new(self.shards.len())
+            .and_then(|len| self.shards[rng % len].lock().take(rng / len))
+    }
+
+    pub fn insert(&self, id: uuid::Uuid, token: CancellationToken) -> CancelGuard<'_> {
+        let shard = NonZeroUsize::new(self.shards.len()).map(|len| {
+            let hash = self.hasher.hash_one(id) as usize;
+            let shard = &self.shards[hash % len];
+            shard.lock().insert(id, token);
+            shard
+        });
+        CancelGuard { shard, id }
+    }
+}
+
+impl CancelShard {
+    fn take(&mut self, rng: usize) -> Option<CancellationToken> {
+        NonZeroUsize::new(self.tokens.len()).and_then(|len| {
+            // 10 second grace period so we don't cancel new connections
+            if self.tokens.get_index(rng % len)?.1 .0.elapsed() < Duration::from_secs(10) {
+                return None;
+            }
+
+            let (_key, (_insert, token)) = self.tokens.swap_remove_index(rng % len)?;
+            Some(token)
+        })
+    }
+
+    fn remove(&mut self, id: uuid::Uuid) {
+        self.tokens.swap_remove(&id);
+    }
+
+    fn insert(&mut self, id: uuid::Uuid, token: CancellationToken) {
+        self.tokens.insert(id, (Instant::now(), token));
+    }
+}
+
+pub struct CancelGuard<'a> {
+    shard: Option<&'a Mutex<CancelShard>>,
+    id: Uuid,
+}
+
+impl Drop for CancelGuard<'_> {
+    fn drop(&mut self) {
+        if let Some(shard) = self.shard {
+            shard.lock().remove(self.id);
+        }
+    }
+}

proxy/src/serverless/conn_pool.rs

@@ -716,7 +716,7 @@ impl<C: ClientInnerExt> Drop for Client<C> {
 mod tests {
     use std::{mem, sync::atomic::AtomicBool};
 
-    use crate::{BranchId, EndpointId, ProjectId};
+    use crate::{serverless::cancel_set::CancelSet, BranchId, EndpointId, ProjectId};
 
     use super::*;
 
@@ -767,6 +767,8 @@ mod tests {
                 max_total_conns: 3,
             },
             request_timeout: Duration::from_secs(1),
+            cancel_set: CancelSet::new(0),
+            client_conn_threshold: u64::MAX,
         }));
         let pool = GlobalConnPool::new(config);
         let conn_info = ConnInfo {

proxy/src/serverless/sql_over_http.rs

@@ -424,8 +424,8 @@ pub enum SqlOverHttpCancel {
 impl ReportableError for SqlOverHttpCancel {
     fn get_error_kind(&self) -> ErrorKind {
         match self {
-            SqlOverHttpCancel::Postgres => ErrorKind::RateLimit,
-            SqlOverHttpCancel::Connect => ErrorKind::ServiceRateLimit,
+            SqlOverHttpCancel::Postgres => ErrorKind::ClientDisconnect,
+            SqlOverHttpCancel::Connect => ErrorKind::ClientDisconnect,
         }
     }
 }