Add request limits for mongo

Author: palfrey

nativelink-config/src/stores.rs

@@ -1441,6 +1441,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")]
@@ -1480,6 +1481,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 {

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, 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,16 @@ 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 % 100 == 0 {
+            debug!(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 +319,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 +338,9 @@ impl StoreDriver for ExperimentalMongoStore {
                     ));
                 }
             }
+            drop(semaphore);
         }
+
         Ok(())
     }
 
@@ -370,6 +395,8 @@ impl StoreDriver for ExperimentalMongoStore {
             }
         }
 
+        let semaphore = self.acquire_permit().await?;
+
         let mut cursor = self
             .cas_collection
             .find(filter)
@@ -394,6 +421,8 @@ impl StoreDriver for ExperimentalMongoStore {
             }
         }
 
+        drop(semaphore);
+
         Ok(count)
     }
 
@@ -463,6 +492,8 @@ impl StoreDriver for ExperimentalMongoStore {
             SIZE_FIELD: size,
         };
 
+        let semaphore = self.acquire_permit().await?;
+
         // Upsert the document
         self.cas_collection
             .update_one(
@@ -473,6 +504,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 +529,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 +588,8 @@ impl StoreDriver for ExperimentalMongoStore {
             }
         }
 
+        drop(semaphore);
+
         writer.send_eof().map_err(|e| {
             make_err!(
                 Code::Internal,
@@ -593,6 +630,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 +957,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 +972,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 +1002,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 +1015,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

@@ -69,6 +69,7 @@ fn create_test_spec_with_key_prefix(key_prefix: Option<String>) -> ExperimentalM
         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,
     }
 }
 
@@ -1074,6 +1075,7 @@ async fn test_non_w_config() -> Result<(), Error> {
             write_concern_w: None,
             write_concern_j: Some(true),
             write_concern_timeout_ms: Some(1),
+            max_request_permits: None,
         })
         .await
         .unwrap_err()