feat: add concurrency budget (#12117)

crates/polars-io/src/parquet/async_impl.rs

@@ -18,7 +18,7 @@ use super::cloud::{build_object_store, CloudLocation, CloudReader};
 use super::mmap::ColumnStore;
 use crate::cloud::CloudOptions;
 use crate::parquet::read_impl::compute_row_group_range;
-use crate::pl_async::get_runtime;
+use crate::pl_async::{get_runtime, with_concurrency_budget};
 use crate::predicates::PhysicalIoExpr;
 use crate::prelude::predicates::read_this_row_group;
 
@@ -50,17 +50,23 @@ impl ParquetObjectStore {
     }
 
     async fn get_range(&self, start: usize, length: usize) -> PolarsResult<Bytes> {
-        self.store
-            .get_range(&self.path, start..start + length)
-            .await
-            .map_err(to_compute_err)
+        with_concurrency_budget(1, || async {
+            self.store
+                .get_range(&self.path, start..start + length)
+                .await
+                .map_err(to_compute_err)
+        })
+        .await
     }
 
     async fn get_ranges(&self, ranges: &[Range<usize>]) -> PolarsResult<Vec<Bytes>> {
-        self.store
-            .get_ranges(&self.path, ranges)
-            .await
-            .map_err(to_compute_err)
+        with_concurrency_budget(ranges.len() as u16, || async {
+            self.store
+                .get_ranges(&self.path, ranges)
+                .await
+                .map_err(to_compute_err)
+        })
+        .await
     }
 
     /// Initialize the length property of the object, unless it has already been fetched.

crates/polars-io/src/pl_async.rs

@@ -1,12 +1,39 @@
 use std::collections::BTreeSet;
 use std::future::Future;
 use std::ops::Deref;
+use std::sync::atomic::{AtomicI32, Ordering};
 use std::sync::RwLock;
+use std::time::Duration;
 
 use once_cell::sync::Lazy;
 use polars_core::POOL;
 use tokio::runtime::{Builder, Runtime};
 
+static CONCURRENCY_BUDGET: AtomicI32 = AtomicI32::new(32);
+
+pub async fn with_concurrency_budget<F, Fut>(requested_budget: u16, callable: F) -> Fut::Output
+where
+    F: FnOnce() -> Fut,
+    Fut: Future,
+{
+    loop {
+        let requested_budget = requested_budget as i32;
+        let available_budget = CONCURRENCY_BUDGET.fetch_sub(requested_budget, Ordering::Relaxed);
+
+        // Bail out, there was no budget
+        if available_budget < 0 {
+            CONCURRENCY_BUDGET.fetch_add(requested_budget, Ordering::Relaxed);
+            tokio::time::sleep(Duration::from_millis(50)).await;
+        } else {
+            let fut = callable();
+            let out = fut.await;
+            CONCURRENCY_BUDGET.fetch_add(requested_budget, Ordering::Relaxed);
+
+            return out;
+        }
+    }
+}
+
 pub struct RuntimeManager {
     rt: Runtime,
     blocking_rayon_threads: RwLock<BTreeSet<usize>>,