Add request limits for mongo

Author: palfrey

nativelink-config/examples/stores-config.json5

@@ -295,7 +295,8 @@
          "key_prefix": "cas:",
          "read_chunk_size": 65536,
          "max_concurrent_uploads": 10,
-         "enable_change_streams": false
+         "enable_change_streams": false,
+         "max_requests": "100"
      }
     },
 ]}

nativelink-config/src/stores.rs

@@ -551,7 +551,8 @@ pub enum StoreSpec {
     ///     "key_prefix": "cas:",
     ///     "read_chunk_size": 65536,
     ///     "max_concurrent_uploads": 10,
-    ///     "enable_change_streams": false
+    ///     "enable_change_streams": false,
+    ///     "max_requests": "100"
     /// }
     /// ```
     ///
@@ -1490,6 +1491,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 +1531,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_requests: Option<usize>,
 }
 
 impl Retry {

nativelink-store/src/mongo_store.rs

@@ -15,6 +15,7 @@
 use core::cmp;
 use core::ops::Bound;
 use core::pin::Pin;
+use core::sync::atomic::{AtomicUsize, Ordering};
 use core::time::Duration;
 use std::borrow::Cow;
 use std::sync::{Arc, Weak};
@@ -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::{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,19 @@ 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 spec.max_concurrent_uploads != 0 {
+            warn!(
+                "max_concurrent_uploads was set for Mongo, and it's a deprecated value we don't use anymore"
+            );
+        }
+
+        if let Some(max_permits) = spec.max_requests {
+            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 +235,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_requests.unwrap_or(Semaphore::MAX_PERMITS),
+            )),
+            waiting_permits: Arc::new(AtomicUsize::new(0)),
         };
 
         Ok(Arc::new(store))
@@ -282,6 +296,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 {
+            info!(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 +328,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 +347,9 @@ impl StoreDriver for ExperimentalMongoStore {
                     ));
                 }
             }
+            drop(semaphore);
         }
+
         Ok(())
     }
 
@@ -370,6 +404,8 @@ impl StoreDriver for ExperimentalMongoStore {
             }
         }
 
+        let semaphore = self.acquire_permit().await?;
+
         let mut cursor = self
             .cas_collection
             .find(filter)
@@ -394,6 +430,8 @@ impl StoreDriver for ExperimentalMongoStore {
             }
         }
 
+        drop(semaphore);
+
         Ok(count)
     }
 
@@ -463,6 +501,8 @@ impl StoreDriver for ExperimentalMongoStore {
             SIZE_FIELD: size,
         };
 
+        let semaphore = self.acquire_permit().await?;
+
         // Upsert the document
         self.cas_collection
             .update_one(
@@ -473,6 +513,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 +538,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 +597,8 @@ impl StoreDriver for ExperimentalMongoStore {
             }
         }
 
+        drop(semaphore);
+
         writer.send_eof().map_err(|e| {
             make_err!(
                 Code::Internal,
@@ -593,6 +639,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 +966,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 +981,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 +1011,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 +1024,8 @@ impl SchedulerStore for ExperimentalMongoStore {
                     make_err!(Code::Internal, "Failed to update scheduler document: {e}")
                 })?;
 
+            drop(semaphore);
+
             Ok(Some(0))
         }
     }

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_requests: 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_requests = 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_requests = 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(())
+}