Add --max-concurrent-requests flag (#1030)

* Add --max-concurrent-requests flag

* Add missing field in test

Author: JakeHedman

libsql-server/src/config.rs

@@ -125,6 +125,7 @@ pub struct DbConfig {
     pub snapshot_exec: Option<String>,
     pub checkpoint_interval: Option<Duration>,
     pub snapshot_at_shutdown: bool,
+    pub max_concurrent_requests: u64,
 }
 
 impl Default for DbConfig {
@@ -141,6 +142,7 @@ impl Default for DbConfig {
             snapshot_exec: None,
             checkpoint_interval: None,
             snapshot_at_shutdown: false,
+            max_concurrent_requests: 128,
         }
     }
 }

libsql-server/src/connection/mod.rs

@@ -159,11 +159,18 @@ pub trait MakeConnection: Send + Sync + 'static {
         semaphore: Arc<Semaphore>,
         timeout: Option<Duration>,
         max_total_response_size: u64,
+        max_concurrent_requests: u64,
     ) -> MakeThrottledConnection<Self>
     where
         Self: Sized,
     {
-        MakeThrottledConnection::new(semaphore, self, timeout, max_total_response_size)
+        MakeThrottledConnection::new(
+            semaphore,
+            self,
+            timeout,
+            max_total_response_size,
+            max_concurrent_requests,
+        )
     }
 }
 
@@ -190,6 +197,7 @@ pub struct MakeThrottledConnection<F> {
     // will result in reducing concurrency to prevent out-of-memory errors.
     max_total_response_size: u64,
     waiters: AtomicUsize,
+    max_concurrent_requests: u64,
 }
 
 impl<F> MakeThrottledConnection<F> {
@@ -198,13 +206,15 @@ impl<F> MakeThrottledConnection<F> {
         connection_maker: F,
         timeout: Option<Duration>,
         max_total_response_size: u64,
+        max_concurrent_requests: u64,
     ) -> Self {
         Self {
             semaphore,
             connection_maker,
             timeout,
             max_total_response_size,
             waiters: AtomicUsize::new(0),
+            max_concurrent_requests,
         }
     }
 
@@ -263,7 +273,7 @@ impl<F: MakeConnection> MakeConnection for MakeThrottledConnection<F> {
         );
         let units = self.units_to_take();
         let waiters_guard = WaitersGuard::new(&self.waiters);
-        if waiters_guard.waiters.load(Ordering::Relaxed) >= 128 {
+        if (waiters_guard.waiters.load(Ordering::Relaxed) as u64) >= self.max_concurrent_requests {
             return Err(Error::TooManyRequests);
         }
         let fut = self.semaphore.clone().acquire_many_owned(units);
@@ -423,6 +433,7 @@ pub mod test {
             Arc::new(Semaphore::new(10)),
             Some(Duration::from_millis(100)),
             u64::MAX,
+            u64::MAX,
         );
 
         let mut conns = Vec::with_capacity(10);

libsql-server/src/lib.rs

@@ -529,6 +529,7 @@ where
             checkpoint_interval: self.db_config.checkpoint_interval,
             disable_namespace: self.disable_namespaces,
             max_concurrent_connections: Arc::new(Semaphore::new(self.max_concurrent_connections)),
+            max_concurrent_requests: self.db_config.max_concurrent_requests,
         };
         let factory = PrimaryNamespaceMaker::new(conf);
         let namespaces = NamespaceStore::new(factory, false, self.db_config.snapshot_at_shutdown);
@@ -619,6 +620,7 @@ impl<C: Connector> Replica<C> {
             max_response_size: self.db_config.max_response_size,
             max_total_response_size: self.db_config.max_total_response_size,
             max_concurrent_connections: Arc::new(Semaphore::new(self.max_concurrent_connections)),
+            max_concurrent_requests: self.db_config.max_concurrent_requests,
         };
         let factory = ReplicaNamespaceMaker::new(conf);
         let namespaces = NamespaceStore::new(factory, true, false);

libsql-server/src/main.rs

@@ -231,6 +231,9 @@ struct Cli {
 
     #[clap(long, default_value = "128", env = "SQLD_MAX_CONCURRENT_CONNECTIONS")]
     max_concurrent_connections: usize,
+    // max number of concurrent requests across all connections
+    #[clap(long, default_value = "128", env = "SQLD_MAX_CONCURRENT_REQUESTS")]
+    max_concurrent_requests: u64,
 }
 
 #[derive(clap::Subcommand, Debug)]
@@ -337,6 +340,7 @@ fn make_db_config(config: &Cli) -> anyhow::Result<DbConfig> {
         snapshot_exec: config.snapshot_exec.clone(),
         checkpoint_interval: config.checkpoint_interval_s.map(Duration::from_secs),
         snapshot_at_shutdown: config.snapshot_at_shutdown,
+        max_concurrent_requests: config.max_concurrent_requests,
     })
 }
 

libsql-server/src/namespace/mod.rs

@@ -669,6 +669,7 @@ pub struct ReplicaNamespaceConfig {
     /// Stats monitor
     pub stats_sender: StatsSender,
     pub max_concurrent_connections: Arc<Semaphore>,
+    pub max_concurrent_requests: u64,
 }
 
 impl Namespace<ReplicaDatabase> {
@@ -789,6 +790,7 @@ impl Namespace<ReplicaDatabase> {
             config.max_concurrent_connections.clone(),
             Some(DB_CREATE_TIMEOUT),
             config.max_total_response_size,
+            config.max_concurrent_requests,
         );
 
         Ok(Self {
@@ -817,6 +819,7 @@ pub struct PrimaryNamespaceConfig {
     pub checkpoint_interval: Option<Duration>,
     pub disable_namespace: bool,
     pub max_concurrent_connections: Arc<Semaphore>,
+    pub max_concurrent_requests: u64,
 }
 
 pub type DumpStream =
@@ -991,6 +994,7 @@ impl Namespace<PrimaryDatabase> {
             config.max_concurrent_connections.clone(),
             Some(DB_CREATE_TIMEOUT),
             config.max_total_response_size,
+            config.max_concurrent_requests,
         )
         .into();
 

libsql-server/src/test/bottomless.rs

@@ -51,6 +51,11 @@ async fn configure_server(
             snapshot_exec: None,
             checkpoint_interval: Some(Duration::from_secs(3)),
             snapshot_at_shutdown: false,
+<<<<<<< HEAD
+=======
+            encryption_key: None,
+            max_concurrent_requests: 128,
+>>>>>>> 63a707798a (Add --max-concurrent-requests flag (#1030))
         },
         admin_api_config: None,
         disable_namespaces: true,