Add request limits for mongo

palfrey · palfrey · commit a794c207eb5a · 2026-03-19T10:42:46.000Z
diff --git a/nativelink-config/src/stores.rs b/nativelink-config/src/stores.rs
@@ -1490,6 +1490,7 @@ pub struct ExperimentalMongoSpec {
     #[serde(default, deserialize_with = "convert_data_size_with_shellexpand")]
     pub read_chunk_size: usize,
 
+    /// Deprecated, unused
     /// Maximum number of concurrent uploads allowed.
     /// Default: 10
     #[serde(default, deserialize_with = "convert_numeric_with_shellexpand")]
@@ -1529,6 +1530,14 @@ pub struct ExperimentalMongoSpec {
         deserialize_with = "convert_optional_numeric_with_shellexpand"
     )]
     pub write_concern_timeout_ms: Option<u32>,
+
+    /// Limits the number of requests at any one time
+    /// Default: Unlimited
+    #[serde(
+        default,
+        deserialize_with = "convert_optional_numeric_with_shellexpand"
+    )]
+    pub max_request_permits: Option<usize>,
 }
 
 impl Retry {
diff --git a/nativelink-error/src/lib.rs b/nativelink-error/src/lib.rs
@@ -353,6 +353,24 @@ impl From<zip::result::ZipError> for Error {
     }
 }
 
+impl From<mongodb::error::Error> for Error {
+    fn from(value: mongodb::error::Error) -> Self {
+        Self::new(Code::Internal, value.to_string())
+    }
+}
+
+impl From<reqwest::Error> for Error {
+    fn from(value: reqwest::Error) -> Self {
+        Self::new(Code::Internal, value.to_string())
+    }
+}
+
+impl From<zip::result::ZipError> for Error {
+    fn from(value: zip::result::ZipError) -> Self {
+        Self::new(Code::Internal, value.to_string())
+    }
+}
+
 pub trait ResultExt<T> {
     /// # Errors
     ///
diff --git a/nativelink-store/src/mongo_store.rs b/nativelink-store/src/mongo_store.rs
@@ -17,6 +17,7 @@ use core::ops::Bound;
 use core::pin::Pin;
 use core::time::Duration;
 use std::borrow::Cow;
+use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::{Arc, Weak};
 
 use async_trait::async_trait;
@@ -39,9 +40,9 @@ use nativelink_util::store_trait::{
 use nativelink_util::task::JoinHandleDropGuard;
 use parking_lot::{Mutex, RwLock};
 use patricia_tree::StringPatriciaMap;
-use tokio::sync::watch;
+use tokio::sync::{Semaphore, SemaphorePermit, watch};
 use tokio::time::sleep;
-use tracing::{error, info, warn};
+use tracing::{debug, error, info, trace, warn};
 
 use crate::cas_utils::is_zero_digest;
 
@@ -63,9 +64,6 @@ const DEFAULT_CONNECTION_TIMEOUT_MS: u64 = 3000;
 /// The default command timeout in milliseconds if not specified.
 const DEFAULT_COMMAND_TIMEOUT_MS: u64 = 10_000;
 
-/// The default maximum number of concurrent uploads.
-const DEFAULT_MAX_CONCURRENT_UPLOADS: usize = 10;
-
 /// The name of the field in `MongoDB` documents that stores the key.
 const KEY_FIELD: &str = "_id";
 
@@ -102,16 +100,18 @@ pub struct ExperimentalMongoStore {
     #[metric(help = "The amount of data to read from MongoDB at a time")]
     read_chunk_size: usize,
 
-    /// The maximum number of concurrent uploads.
-    #[metric(help = "The maximum number of concurrent uploads")]
-    max_concurrent_uploads: usize,
-
     /// Enable change streams for real-time updates.
     #[metric(help = "Whether change streams are enabled")]
     enable_change_streams: bool,
 
     /// A manager for subscriptions to keys in `MongoDB`.
     subscription_manager: Mutex<Option<Arc<ExperimentalMongoSubscriptionManager>>>,
+
+    /// Limits the number of requests at any one time
+    request_permits: Arc<Semaphore>,
+
+    /// Keep track of the request_permits queue size
+    waiting_permits: Arc<AtomicUsize>,
 }
 
 impl ExperimentalMongoStore {
@@ -149,8 +149,13 @@ impl ExperimentalMongoStore {
             spec.command_timeout_ms = DEFAULT_COMMAND_TIMEOUT_MS;
         }
 
-        if spec.max_concurrent_uploads == 0 {
-            spec.max_concurrent_uploads = DEFAULT_MAX_CONCURRENT_UPLOADS;
+        if let Some(max_permits) = spec.max_request_permits {
+            if max_permits == 0 {
+                return Err(make_err!(
+                    Code::InvalidArgument,
+                    "max_request_permits was set to zero, which will block mongo_store from working at all"
+                ));
+            }
         }
 
         // Configure client options
@@ -224,9 +229,12 @@ impl ExperimentalMongoStore {
             scheduler_collection,
             key_prefix: spec.key_prefix.clone().unwrap_or_default(),
             read_chunk_size: spec.read_chunk_size,
-            max_concurrent_uploads: spec.max_concurrent_uploads,
             enable_change_streams: spec.enable_change_streams,
             subscription_manager: Mutex::new(None),
+            request_permits: Arc::new(Semaphore::new(
+                spec.max_request_permits.unwrap_or(Semaphore::MAX_PERMITS),
+            )),
+            waiting_permits: Arc::new(AtomicUsize::new(0)),
         };
 
         Ok(Arc::new(store))
@@ -282,6 +290,19 @@ impl ExperimentalMongoStore {
             key.strip_prefix(&self.key_prefix).map(ToString::to_string)
         }
     }
+
+    async fn acquire_permit(&self) -> Result<SemaphorePermit<'_>, Error> {
+        let waiting = self.waiting_permits.fetch_add(1, Ordering::Relaxed);
+
+        if waiting > 0 && waiting % 100 == 0 {
+            debug!(waiting, "Number of waiting permits for Mongo");
+        } else {
+            trace!(waiting, "Number of waiting permits for Mongo");
+        }
+        let permit = self.request_permits.acquire().await;
+        self.waiting_permits.fetch_sub(1, Ordering::Relaxed);
+        Ok(permit?)
+    }
 }
 
 #[async_trait]
@@ -301,6 +322,11 @@ impl StoreDriver for ExperimentalMongoStore {
             let encoded_key = self.encode_key(key);
             let filter = doc! { KEY_FIELD: encoded_key.as_ref() };
 
+            // We could do this with acquire_many, but that's unsafe if the number of keys is greater
+            // than the number of permits, as it'll block forever. Doing this one at a time is guaranteed
+            // not to block provided no-one sets permits to 0, and we check for that case at startup.
+            let semaphore = self.acquire_permit().await?;
+
             match self.cas_collection.find_one(filter).await {
                 Ok(Some(doc)) => {
                     *result = doc.get_i64(SIZE_FIELD).ok().map(|v| v as u64);
@@ -315,7 +341,9 @@ impl StoreDriver for ExperimentalMongoStore {
                     ));
                 }
             }
+            drop(semaphore);
         }
+
         Ok(())
     }
 
@@ -370,6 +398,8 @@ impl StoreDriver for ExperimentalMongoStore {
             }
         }
 
+        let semaphore = self.acquire_permit().await?;
+
         let mut cursor = self
             .cas_collection
             .find(filter)
@@ -394,6 +424,8 @@ impl StoreDriver for ExperimentalMongoStore {
             }
         }
 
+        drop(semaphore);
+
         Ok(count)
     }
 
@@ -463,6 +495,8 @@ impl StoreDriver for ExperimentalMongoStore {
             SIZE_FIELD: size,
         };
 
+        let semaphore = self.acquire_permit().await?;
+
         // Upsert the document
         self.cas_collection
             .update_one(
@@ -473,6 +507,8 @@ impl StoreDriver for ExperimentalMongoStore {
             .await
             .map_err(|e| make_err!(Code::Internal, "Failed to update document in MongoDB: {e}"))?;
 
+        drop(semaphore);
+
         Ok(())
     }
 
@@ -496,6 +532,8 @@ impl StoreDriver for ExperimentalMongoStore {
         let encoded_key = self.encode_key(&key);
         let filter = doc! { KEY_FIELD: encoded_key.as_ref() };
 
+        let semaphore = self.acquire_permit().await?;
+
         let doc = self
             .cas_collection
             .find_one(filter)
@@ -553,6 +591,8 @@ impl StoreDriver for ExperimentalMongoStore {
             }
         }
 
+        drop(semaphore);
+
         writer.send_eof().map_err(|e| {
             make_err!(
                 Code::Internal,
@@ -593,6 +633,8 @@ impl HealthStatusIndicator for ExperimentalMongoStore {
     }
 
     async fn check_health(&self, namespace: Cow<'static, str>) -> HealthStatus {
+        // Note we do not acquire a request_permit here, as the health check needs to always go through
+        // even if everything else is fully loaded
         match self.database.run_command(doc! { "ping": 1 }).await {
             Ok(_) => HealthStatus::new_ok(self, "Connection healthy".into()),
             Err(e) => HealthStatus::new_failed(
@@ -918,7 +960,9 @@ impl SchedulerStore for ExperimentalMongoStore {
                 VERSION_FIELD: current_version,
             };
 
-            match self
+            let semaphore = self.acquire_permit().await?;
+
+            let result = match self
                 .scheduler_collection
                 .find_one_and_update(filter, update_doc)
                 .upsert(true)
@@ -931,7 +975,9 @@ impl SchedulerStore for ExperimentalMongoStore {
                     Code::Internal,
                     "MongoDB error in update_data: {e}"
                 )),
-            }
+            };
+            drop(semaphore);
+            result
         } else {
             let data_bytes = data.try_into_bytes().map_err(|e| {
                 make_err!(
@@ -959,6 +1005,8 @@ impl SchedulerStore for ExperimentalMongoStore {
                 );
             }
 
+            let semaphore = self.acquire_permit().await?;
+
             self.scheduler_collection
                 .update_one(
                     doc! { KEY_FIELD: encoded_key.as_ref() },
@@ -970,6 +1018,8 @@ impl SchedulerStore for ExperimentalMongoStore {
                     make_err!(Code::Internal, "Failed to update scheduler document: {e}")
                 })?;
 
+            drop(semaphore);
+
             Ok(Some(0))
         }
     }
diff --git a/nativelink-store/tests/mongo_store_test.rs b/nativelink-store/tests/mongo_store_test.rs
@@ -72,6 +72,7 @@ fn create_test_spec_with_key_prefix(
         write_concern_w: Some("majority".to_string()),
         write_concern_j: Some(true),
         write_concern_timeout_ms: Some(10000), // Longer timeout for remote DB
+        max_request_permits: None,
     }
 }
 
@@ -967,25 +968,55 @@ async fn test_scheduler_store_operations() -> Result<(), Error> {
 
 #[nativelink_test]
 async fn test_non_w_config() -> Result<(), Error> {
+    let (mut spec, mongo_process) = TestMongoHelper::new_spec(None).await?;
+    spec.write_concern_w = None;
+    spec.write_concern_j = Some(true);
+    spec.write_concern_timeout_ms = Some(1);
+
     assert_eq!(
         Error::new(Code::InvalidArgument, "write_concern_w not set, but j and/or timeout set. Please set 'write_concern_w' to a non-default value. See https://www.mongodb.com/docs/manual/reference/write-concern/#w-option for options.".to_string()),
-        ExperimentalMongoStore::new(ExperimentalMongoSpec {
-            connection_string: "mongodb://dummy".to_string(),
-            database: "dummy".to_string(),
-            cas_collection: "test_cas".to_string(),
-            scheduler_collection: "test_scheduler".to_string(),
-            key_prefix: None,
-            read_chunk_size: 1024,
-            max_concurrent_uploads: 10,
-            connection_timeout_ms: 10000,
-            command_timeout_ms: 15000,
-            enable_change_streams: false,
-            write_concern_w: None,
-            write_concern_j: Some(true),
-            write_concern_timeout_ms: Some(1),
-        })
-        .await
+        TestMongoHelper::new_with_spec_and_process(spec, mongo_process).await
         .unwrap_err()
     );
     Ok(())
 }
+
+#[nativelink_test]
+async fn empty_request_permit() -> Result<(), Error> {
+    let (mut spec, mongo_process) = TestMongoHelper::new_spec(None).await?;
+    spec.max_request_permits = Some(0);
+
+    assert_eq!(
+        Error::new(
+            Code::InvalidArgument,
+            "max_request_permits was set to zero, which will block mongo_store from working at all"
+                .to_string()
+        ),
+        TestMongoHelper::new_with_spec_and_process(spec, mongo_process)
+            .await
+            .unwrap_err()
+    );
+    Ok(())
+}
+
+#[nativelink_test]
+async fn single_request_permit() -> Result<(), Error> {
+    let (mut spec, mongo_process) = TestMongoHelper::new_spec(None).await?;
+    spec.max_request_permits = Some(1);
+    let helper = TestMongoHelper::new_with_spec_and_process(spec, mongo_process).await?;
+
+    let data = Bytes::from_static(b"14");
+    let digest = DigestInfo::try_new(VALID_HASH1, 2)?;
+    helper.store.update_oneshot(digest, data.clone()).await?;
+    let result = helper.store.has(digest).await?;
+    assert!(
+        result.is_some(),
+        "Expected mongo store to have hash: {VALID_HASH1}",
+    );
+
+    assert!(logs_contain(
+        "Number of waiting permits for Mongo waiting=0"
+    ));
+
+    Ok(())
+}
