Bound object store reads with an adaptive permit pool, Azure first

nativelink-config/src/stores.rs

@@ -1433,6 +1433,44 @@ pub struct CommonObjectSpec {
     /// Default: 120.0
     #[serde(default = "default_throttle_retry_budget_s")]
     pub throttle_retry_budget_s: f32,
+
+    /// Most downloads (`GET`) this process keeps in flight against the
+    /// object store at once, each counted until its body has been sent on.
+    /// Further downloads wait in the process for a permit instead of being
+    /// sent and refused, which turns an account over its limit into
+    /// queueing here rather than wasted requests and retries in lockstep.
+    /// Existence checks (`HEAD`) are not counted: they move no bytes, so
+    /// they never wait behind a body going to a slow client, and a
+    /// throttled `HEAD` does not move the limit. 0 (the default) leaves
+    /// downloads unbounded, as before.
+    ///
+    /// Default: 0
+    #[serde(default, deserialize_with = "convert_numeric_with_shellexpand")]
+    pub max_concurrent_reads: usize,
+
+    /// With `max_concurrent_reads` set, let the bound move: a throttled
+    /// download (503 `ServerBusy` or 429) halves it, not below
+    /// `min_concurrent_reads`, and a run of answered downloads raises it
+    /// by one back toward `max_concurrent_reads`. Many processes behind one
+    /// account then converge on a share of its limit. Off, the bound is
+    /// fixed at `max_concurrent_reads`.
+    ///
+    /// Default: false
+    #[serde(default, deserialize_with = "convert_boolean_with_shellexpand")]
+    pub adaptive_read_concurrency: bool,
+
+    /// Floor for the adaptive read bound.
+    ///
+    /// Default: 4
+    #[serde(
+        default = "default_min_concurrent_reads",
+        deserialize_with = "convert_numeric_with_shellexpand"
+    )]
+    pub min_concurrent_reads: usize,
+}
+
+const fn default_min_concurrent_reads() -> usize {
+    4
 }
 
 const fn default_throttle_retry_delay_s() -> f32 {
@@ -1460,6 +1498,9 @@ impl Default for CommonObjectSpec {
             throttle_retry_delay_s: default_throttle_retry_delay_s(),
             throttle_retry_max_delay_s: default_throttle_retry_max_delay_s(),
             throttle_retry_budget_s: default_throttle_retry_budget_s(),
+            max_concurrent_reads: 0,
+            adaptive_read_concurrency: false,
+            min_concurrent_reads: default_min_concurrent_reads(),
         }
     }
 }

nativelink-store/src/azure_blob_store.rs

@@ -35,6 +35,7 @@ use futures::{FutureExt, StreamExt, TryStreamExt};
 use nativelink_config::stores::ExperimentalAzureSpec;
 use nativelink_error::{Code, Error, ResultExt, make_err};
 use nativelink_metric::MetricsComponent;
+use nativelink_util::adaptive_permits::{AdaptivePermit, AdaptivePermits};
 use nativelink_util::buf_channel::{
     DropCloserReadHalf, DropCloserWriteHalf, make_buf_channel_pair,
 };
@@ -74,6 +75,9 @@ const DEFAULT_MAX_CONCURRENT_UPLOADS: usize = 10;
 // Default public Azure Blob Storage endpoint suffix.
 const DEFAULT_BLOB_ENDPOINT_SUFFIX: &str = "blob.core.windows.net";
 
+/// Successful reads in a row before the adaptive read limit grows by one.
+const READ_LIMIT_GROW_AFTER: usize = 16;
+
 #[derive(MetricsComponent)]
 pub struct AzureBlobStore<NowFn> {
     client: Arc<BlobContainerClient>,
@@ -91,6 +95,9 @@ pub struct AzureBlobStore<NowFn> {
     max_concurrent_uploads: usize,
     throttle_retry_delay: Duration,
     throttle_retry_max_delay: Duration,
+    /// Bounds `GET` and `HEAD` requests in flight; `None` leaves them
+    /// unbounded. See `CommonObjectSpec::max_concurrent_reads`.
+    read_permits: Option<Arc<AdaptivePermits>>,
 }
 
 impl<NowFn> core::fmt::Debug for AzureBlobStore<NowFn> {
@@ -204,6 +211,18 @@ where
                 spec.common.throttle_retry_delay_s
             ));
         }
+        let read_permits = (spec.common.max_concurrent_reads > 0).then(|| {
+            let max = spec.common.max_concurrent_reads;
+            let (min, grow_after) = if spec.common.adaptive_read_concurrency {
+                (
+                    spec.common.min_concurrent_reads.min(max),
+                    READ_LIMIT_GROW_AFTER,
+                )
+            } else {
+                (max, 0)
+            };
+            AdaptivePermits::new("azure", max, min, max, grow_after)
+        });
         let throttle_retry_budget =
             Duration::try_from_secs_f32(spec.common.throttle_retry_budget_s).map_err(|e| {
                 make_err!(
@@ -240,9 +259,42 @@ where
                 .unwrap_or(DEFAULT_MAX_CONCURRENT_UPLOADS),
             throttle_retry_delay,
             throttle_retry_max_delay,
+            read_permits,
         }))
     }
 
+    /// The read limit the permit pool stands at, or `None` with reads
+    /// unbounded.
+    pub fn read_limit(&self) -> Option<usize> {
+        self.read_permits.as_ref().map(|pool| pool.limit())
+    }
+
+    /// A download permit when reads are bounded. Held for the request,
+    /// body included, so what is in flight is what the limit says. Only
+    /// `get_part` takes one: a `HEAD` moves no bytes and must not wait
+    /// behind a body still being sent to a slow client, and a throttle
+    /// on either side says nothing about the other.
+    async fn read_permit(&self) -> Option<AdaptivePermit> {
+        match &self.read_permits {
+            Some(pool) => Some(pool.acquire().await),
+            None => None,
+        }
+    }
+
+    /// A read was answered (found or not): one toward growing the limit.
+    fn note_read_ok(&self) {
+        if let Some(pool) = &self.read_permits {
+            pool.record_success();
+        }
+    }
+
+    /// A read was throttled: the limit is cut.
+    fn note_read_throttled(&self) {
+        if let Some(pool) = &self.read_permits {
+            pool.record_throttle();
+        }
+    }
+
     fn make_blob_path(&self, key: &StoreKey<'_>) -> String {
         format!("{}{}", self.blob_prefix, key.as_str())
     }
@@ -781,6 +833,9 @@ where
                 let client = Arc::clone(&self.client);
                 let blob_path = blob_path.clone();
                 async move {
+                    // Held through the body: what is in flight is what the
+                    // read limit says, not only what has been started.
+                    let _read_permit = self.read_permit().await;
                     let _permit = match fs::get_permit().await {
                         Ok(permit) => permit,
                         Err(e) => {
@@ -875,11 +930,18 @@ where
                     .await;
 
                     match result {
-                        Ok(()) => Some((RetryResult::Ok(()), writer)),
+                        Ok(()) => {
+                            self.note_read_ok();
+                            Some((RetryResult::Ok(()), writer))
+                        }
                         Err((e, _)) if e.code == Code::NotFound => {
+                            self.note_read_ok();
                             Some((RetryResult::Err(e), writer))
                         }
-                        Err((e, Some(wait))) => Some((RetryResult::RetryAfter(e, wait), writer)),
+                        Err((e, Some(wait))) => {
+                            self.note_read_throttled();
+                            Some((RetryResult::RetryAfter(e, wait), writer))
+                        }
                         Err((e, None)) => Some((RetryResult::Retry(e), writer)),
                     }
                 }

nativelink-store/tests/azure_blob_store_test.rs

@@ -200,6 +200,20 @@ fn create_test_store(
     mock: Arc<MockTransport>,
     retry: Option<Retry>,
     consider_expired_after_s: Option<u32>,
+) -> Result<TestStore, Error> {
+    create_test_store_with_common(
+        mock,
+        CommonObjectSpec {
+            consider_expired_after_s: consider_expired_after_s.unwrap_or(0),
+            retry: retry.unwrap_or_default(),
+            ..Default::default()
+        },
+    )
+}
+
+fn create_test_store_with_common(
+    mock: Arc<MockTransport>,
+    common: CommonObjectSpec,
 ) -> Result<TestStore, Error> {
     let mut options = BlobContainerClientOptions::default();
     // NativeLink's own `Retrier` wraps each operation, so disable the SDK's retry
@@ -217,11 +231,7 @@ fn create_test_store(
     let spec = ExperimentalAzureSpec {
         account_name: TEST_ACCOUNT.to_string(),
         container: TEST_CONTAINER.to_string(),
-        common: CommonObjectSpec {
-            consider_expired_after_s: consider_expired_after_s.unwrap_or(0),
-            retry: retry.unwrap_or_default(),
-            ..Default::default()
-        },
+        common,
         ..Default::default()
     };
 
@@ -890,3 +900,126 @@ async fn the_requested_range_starts_at_the_offset_plus_what_was_sent() -> Result
     }
     Ok(())
 }
+
+/// With adaptive read concurrency on, each throttled response halves the
+/// read limit (the two here are a retry wait apart, so they are two
+/// events), and a run of answered reads grows it back one at a time.
+#[nativelink_test(start_paused = true)]
+async fn throttled_reads_halve_the_read_limit_and_successes_regrow_it() -> Result<(), Error> {
+    const VALUE: &str = "test_content";
+
+    let mock = MockTransport::new();
+    mock.push_download(error_response(StatusCode::ServiceUnavailable));
+    mock.push_download(error_response(StatusCode::ServiceUnavailable));
+    mock.push_download(download_ok(StatusCode::Ok, VALUE.as_bytes().to_vec()));
+
+    let store = create_test_store_with_common(
+        Arc::clone(&mock),
+        CommonObjectSpec {
+            retry: Retry {
+                max_retries: 0,
+                delay: 0.0,
+                jitter: 0.0,
+                ..Default::default()
+            },
+            max_concurrent_reads: 16,
+            adaptive_read_concurrency: true,
+            min_concurrent_reads: 2,
+            ..Default::default()
+        },
+    )?;
+    assert_eq!(store.read_limit(), Some(16));
+
+    let result = store
+        .get_part_unchunked(create_test_digest()?, 0, None)
+        .await?;
+    assert_eq!(result, VALUE.as_bytes());
+    assert_eq!(mock.request_count(), 3);
+    assert_eq!(
+        store.read_limit(),
+        Some(4),
+        "two throttles, two halvings; the download that got through is one toward regrowth"
+    );
+
+    // Sixteen answered downloads in a row (the one above was the first)
+    // add one permit.
+    for _ in 0..15 {
+        mock.push_download(download_ok(StatusCode::Ok, VALUE.as_bytes().to_vec()));
+        store
+            .get_part_unchunked(create_test_digest()?, 0, None)
+            .await?;
+    }
+    assert_eq!(store.read_limit(), Some(5));
+    Ok(())
+}
+
+/// Existence checks are outside the download pool: a throttled `HEAD`
+/// does not cut the limit and an answered one does not grow it.
+#[nativelink_test(start_paused = true)]
+async fn existence_checks_do_not_move_the_read_limit() -> Result<(), Error> {
+    let mock = MockTransport::new();
+    mock.push_properties(error_response(StatusCode::ServiceUnavailable));
+    mock.push_properties(properties_ok(512));
+
+    let store = create_test_store_with_common(
+        Arc::clone(&mock),
+        CommonObjectSpec {
+            retry: Retry {
+                max_retries: 0,
+                delay: 0.0,
+                jitter: 0.0,
+                ..Default::default()
+            },
+            max_concurrent_reads: 16,
+            adaptive_read_concurrency: true,
+            min_concurrent_reads: 2,
+            ..Default::default()
+        },
+    )?;
+
+    assert_eq!(store.has(create_test_digest()?).await?, Some(512));
+    assert_eq!(mock.request_count(), 2, "the throttled HEAD was retried");
+    assert_eq!(
+        store.read_limit(),
+        Some(16),
+        "a throttled HEAD cuts nothing"
+    );
+
+    for _ in 0..16 {
+        mock.push_properties(properties_ok(512));
+        store.has(create_test_digest()?).await?;
+    }
+    assert_eq!(store.read_limit(), Some(16), "answered HEADs grow nothing");
+    Ok(())
+}
+
+/// Without `adaptive_read_concurrency` the bound is fixed at
+/// `max_concurrent_reads` whatever the store answers.
+#[nativelink_test(start_paused = true)]
+async fn a_fixed_read_limit_does_not_move_on_a_throttle() -> Result<(), Error> {
+    const VALUE: &str = "test_content";
+
+    let mock = MockTransport::new();
+    mock.push_download(error_response(StatusCode::ServiceUnavailable));
+    mock.push_download(download_ok(StatusCode::Ok, VALUE.as_bytes().to_vec()));
+
+    let store = create_test_store_with_common(
+        Arc::clone(&mock),
+        CommonObjectSpec {
+            retry: Retry {
+                max_retries: 0,
+                delay: 0.0,
+                jitter: 0.0,
+                ..Default::default()
+            },
+            max_concurrent_reads: 8,
+            ..Default::default()
+        },
+    )?;
+
+    store
+        .get_part_unchunked(create_test_digest()?, 0, None)
+        .await?;
+    assert_eq!(store.read_limit(), Some(8));
+    Ok(())
+}

nativelink-util/BUILD.bazel

@@ -14,6 +14,7 @@ rust_library(
     name = "nativelink-util",
     srcs = [
         "src/action_messages.rs",
+        "src/adaptive_permits.rs",
         "src/buf_channel.rs",
         "src/channel_body_for_tests.rs",
         "src/chunked_stream.rs",
@@ -100,6 +101,7 @@ rust_test_suite(
     timeout = "short",
     srcs = [
         "tests/action_messages_test.rs",
+        "tests/adaptive_permits_test.rs",
         "tests/buf_channel_test.rs",
         "tests/cache_size_metrics_test.rs",
         "tests/channel_body_for_tests_test.rs",