polars_io/cloud/concurrency/
mod.rs

1//! Adaptive in-flight concurrency controller for cloud IO (Input only).
2//!
3//! Admission control for concurrency uses two budgets:
4//! - A (primary) bytes-based budget to model the bandwidth-delay product (BDP)
5//! - A (secondary) count-based budget to limit the number of in-flight requests
6//!
7//! The bytes-based budget models the BDP as
8//!   BDP = BW_max * TTFB_est
9//!
10//! Three components cooperate:
11//! - Model: records IO observations and models the network (BW_max, TTFB_est, BDP)
12//! - Regime: state machine driving the admission (Init / RampUp / Stable / ProbeUp)
13//! - Admission: admission control, enforces byte-based + request-based budgets
14
15// Loosely based on BBR: Congestion-Based Congestion Control
16// see https://queue.acm.org/detail.cfm?id=3022184
17
18mod admission;
19mod model;
20mod regime;
21
22use std::num::NonZeroU64;
23use std::sync::Arc;
24use std::time::{Duration, Instant};
25
26pub use admission::{InFlightBudget, InFlightPermit, InFlightStats};
27use crossbeam_queue::ArrayQueue;
28pub use model::Model;
29use polars_core::runtime::ASYNC;
30use polars_utils::relaxed_cell::RelaxedCell;
31pub use regime::{Regime, RegimeState};
32
33// Number of samples in the queue, which gets drained on every tick.
34// At 50k requests per second and 100 ms tick window, we need 5k.
35const SAMPLE_QUEUE_CAPACITY: usize = 8192;
36
37use crate::cloud::concurrency_config::get_random_access_chunk_size;
38use crate::cloud::http_rate_limit::PacingBudget;
39
40#[derive(Clone, Copy, Debug)]
41pub struct IoSample {
42    pub n_bytes: u64,
43    // Time-to-first-byte.
44    pub ttfb: Duration,
45    // TODO: Factor out as we only care about per-tick_window stats.
46    pub completion_time: Instant,
47}
48
49#[derive(Debug, Clone)]
50pub struct ControllerConfig {
51    // Sliding window over which the most recent round-trip-time (RTT) and bandwidth (BW)
52    // will be calculated. Also acts as the retention window.
53    window: Duration,
54    // Byte-based budget during the Init phase, and as the base for the RampUp phase.
55    init_byte_budget: u64,
56    // Lower limit for the byte-based budget - needed to avoid deadlock.
57    floor_byte_budget: u64,
58    // Count-based request budget.
59    request_budget: u64,
60    // Lower limit for the count-based budget.
61    floor_request_budget: u64,
62    // Controller regime update frequency.
63    control_interval: Duration,
64    // Total budget only resizes if the relative changes exceeds this threshold
65    budget_resize_threshold: f64,
66}
67
68impl Default for ControllerConfig {
69    fn default() -> Self {
70        // Only used for bytes-based budget.
71        let target_chunk_size = get_random_access_chunk_size() as u64;
72        Self {
73            window: Duration::from_millis(1000),
74
75            // Byte-based budget during the ramp-up phase.
76            // Starting too low results in lost opportunity (time) during ramp-up.
77            // Starting too high results in early congestion, delayed completion of the first chunk,
78            // and inflated bandwidth estimation.
79            //
80            // Some BDP numbers for reference:
81            //   1 Gbps x 20 ms = 2.5 MB
82            //   1 Gbps x 50 ms = 6.25 MB
83            //   10 Gbps x 50 ms = 62.5 MB
84            //   100 Gbps x 50 ms = 625 MB
85            init_byte_budget: get_init_byte_budget(target_chunk_size),
86
87            // Byte-based budget floor.
88            // Must be >=larger than target_chunk_size to avoid potential deadlock.
89            floor_byte_budget: target_chunk_size,
90
91            // Count-based budget.
92            request_budget: get_request_budget(),
93            floor_request_budget: get_floor_request_budget(),
94            control_interval: Duration::from_millis(100),
95            budget_resize_threshold: 0.05,
96        }
97    }
98}
99
100/// Max number of bytes concurrently in flight during the init and start of rampup phase.
101fn get_init_byte_budget(target_chunk_size: u64) -> u64 {
102    let init_byte_budget = std::env::var("POLARS_INFLIGHT_INIT_BYTE_BUDGET")
103        .map(|x| {
104            x.parse::<NonZeroU64>()
105                .unwrap_or_else(|_| {
106                    panic!("invalid value for POLARS_INFLIGHT_INIT_BYTE_BUDGET: {x}")
107                })
108                .get()
109        })
110        .unwrap_or_else(|_| {
111            // This should be lower than the expected BDP so it can ramp-up, but
112            // too low a value delays the transition to stable.
113            // Heuristic: higher bandwidth is expected on larger instances.
114            let n = polars_config::config().max_threads() as u64;
115            n.div_ceil(8).max(4) * target_chunk_size
116        })
117        .max(1);
118
119    if init_byte_budget < target_chunk_size {
120        panic!("in-flight byte budget init must be larger than the target_chunk_size");
121    }
122
123    init_byte_budget
124}
125
126/// Maximum number of requests concurrently in flight.
127pub fn get_request_budget() -> u64 {
128    // Since object_store/reqwest use HTTP/1 with a connection pool, this value controls the
129    // max concurrent TCP sessions to S3 for the pipeline.
130    // When modifying this value, consider the max_thread count configuration(s), the OS limitations
131    // (e.g., ulimit -n), and any cloud infrastructure limitations.
132    std::env::var("POLARS_INFLIGHT_REQUEST_BUDGET")
133        .map(|x| {
134            x.parse::<NonZeroU64>()
135                .unwrap_or_else(|_| panic!("invalid value for POLARS_INFLIGHT_REQUEST_BUDGET: {x}"))
136                .get()
137        })
138        .unwrap_or(512)
139        .max(1)
140}
141
142/// Minimum number of requests concurrently in flight, if demand is there.
143pub fn get_floor_request_budget() -> u64 {
144    // Since object_store/reqwest use HTTP/1 with a connection pool, this value controls the
145    // max concurrent TCP sessions to S3 for the pipeline.
146    // When modifying this value, consider the max_thread count configuration(s), the OS limitations
147    // (e.g., ulimit -n), and any cloud infrastructure limitations.
148    std::env::var("POLARS_INFLIGHT_FLOOR_REQUEST_BUDGET")
149        .map(|x| {
150            x.parse::<NonZeroU64>()
151                .unwrap_or_else(|_| {
152                    panic!("invalid value for POLARS_INFLIGHT_FLOOR_REQUEST_BUDGET: {x}")
153                })
154                .get()
155        })
156        .unwrap_or(polars_config::config().max_threads() as u64)
157        .max(1)
158}
159
160#[derive(Debug)]
161pub struct ConcurrencyController {
162    config: ControllerConfig,
163    sample_queue: Arc<ArrayQueue<IoSample>>,
164    samples_dropped: Arc<RelaxedCell<u64>>,
165    inflight_budget: Arc<InFlightBudget>,
166    _control_task: tokio::task::JoinHandle<()>,
167}
168
169impl ConcurrencyController {
170    pub fn new(config: ControllerConfig, pacing_budget: Option<PacingBudget>) -> Self {
171        let sample_queue = Arc::new(ArrayQueue::new(SAMPLE_QUEUE_CAPACITY));
172        let samples_dropped = Arc::new(RelaxedCell::new_u64(0));
173
174        let inflight_budget = Arc::new(InFlightBudget::new(
175            config.init_byte_budget,
176            config.floor_byte_budget,
177            config.request_budget,
178            config.floor_request_budget,
179        ));
180
181        let control_task = Self::spawn_control_loop(
182            sample_queue.clone(),
183            samples_dropped.clone(),
184            inflight_budget.clone(),
185            config.clone(),
186            pacing_budget,
187        );
188
189        Self {
190            config,
191            sample_queue,
192            samples_dropped,
193            inflight_budget,
194            _control_task: control_task,
195        }
196    }
197
198    pub fn config(&self) -> &ControllerConfig {
199        &self.config
200    }
201
202    /// Record a completed IO. Hot path.
203    pub fn record_io(&self, sample: IoSample) {
204        if self.sample_queue.push(sample).is_err() {
205            // Queue full: drop. Samples are statistics is considered acceptable.
206            self.samples_dropped.fetch_add(1);
207        }
208    }
209
210    pub fn inflight_budget(&self) -> &Arc<InFlightBudget> {
211        &self.inflight_budget
212    }
213
214    pub async fn acquire(&self, bytes: u64) -> InFlightPermit {
215        self.inflight_budget.acquire(bytes).await
216    }
217
218    fn spawn_control_loop(
219        sample_queue: Arc<ArrayQueue<IoSample>>,
220        samples_dropped: Arc<RelaxedCell<u64>>,
221        admission: Arc<InFlightBudget>,
222        config: ControllerConfig,
223        pacing_budget: Option<PacingBudget>,
224    ) -> tokio::task::JoinHandle<()> {
225        if polars_config::config().verbose() {
226            eprintln!(
227                "[InFlightConcurrency]: spawn control loop: control_interval: {}ms",
228                config.control_interval.as_millis()
229            );
230        }
231        ASYNC.spawn(async move {
232            let mut model = Model::new(config.window);
233            let mut regime = Regime::new(Instant::now());
234
235            let mut ticker = tokio::time::interval(config.control_interval);
236            ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
237
238            loop {
239                ticker.tick().await;
240                let now = Instant::now();
241
242                // Limit concurrency to the pacing budget from the rate-limiter.
243                if let Some(ref pacing_budget) = pacing_budget {
244                    let rate = pacing_budget.rate();
245                    let horizon_s = pacing_budget.horizon().as_secs_f64();
246                    let request_budget = rate * horizon_s;
247                    admission.resize_request_budget(request_budget as u64);
248                }
249
250                // Update model statistics and step regime.
251                let (state, signal, dropped, bw_hwm_held) = {
252                    for _ in 0..SAMPLE_QUEUE_CAPACITY {
253                        let Some(s) = sample_queue.pop() else { break };
254                        model.record(s);
255                    }
256                    let dropped = samples_dropped.swap(0);
257                    model.update(now);
258                    let signal = model.signal();
259                    let state = regime.step(signal, now);
260                    let bw_hwm_held = model.bw_hwm_bps();
261                    (state, signal, dropped, bw_hwm_held)
262                };
263
264                if !matches!(state, RegimeState::WarmIdle { .. }) {
265                    // Compute base BDP
266                    let base_byte_budget = match (state, signal) {
267                        (RegimeState::Init, _) | (_, None) => config.init_byte_budget,
268                        (_, Some(signal)) => signal.bdp_bytes().max(config.init_byte_budget),
269                    };
270
271                    // Compute target BDP using the gain multiplier. This is similar to BBR cwnd_gain.
272                    let gain = match state {
273                        RegimeState::Init => 1.0,
274                        RegimeState::RampUp { .. } => 2.0,
275                        // NOTE: >> 1.0 for the purpose of absorbing environment noise.
276                        RegimeState::Stable => 2.0,
277                        RegimeState::ProbeUp { .. } => 3.0,
278                        // Unreachable.
279                        RegimeState::WarmIdle { .. } => 1.0,
280                    };
281                    let target_byte_budget = (base_byte_budget as f64 * gain) as u64;
282
283                    // Resize if needed.
284                    let current_byte_budget = admission.current_byte_budget();
285                    let threshold = config.budget_resize_threshold;
286                    let should_resize = match current_byte_budget {
287                        0 => target_byte_budget > 0,
288                        current => {
289                            let ratio = target_byte_budget as f64 / current as f64;
290                            ratio < (1.0 - threshold) || ratio > (1.0 + threshold)
291                        },
292                    };
293
294                    if should_resize {
295                        admission.resize_byte_budget(target_byte_budget);
296                    }
297                }
298
299                // Log snapshot.
300                if std::env::var("POLARS_LOG_CONCURRENCY").is_ok() {
301                    let stats = admission.stats();
302                    eprintln!(
303                        "[InFlightConcurrency {}] regime={}, \
304                        bw_hwm={:.1} MB/s, \
305                        bw_avg={:.1} MB/s, \
306                        rtt_min={:.1} ms, \
307                        rtt_avg={:.1} ms, \
308                        bdp_obs={:.1} MB, \
309                        bytes_budget={:.1} MB, \
310                        bytes_in_use={:.1} MB, \
311                        bytes_sat={:.2}, \
312                        req_budget={}, \
313                        req_in_use={}, \
314                        req_sat={:.2}",
315                        chrono::Utc::now(),
316                        state.label(),
317                        signal.map(|s| s.bw_hwm_bps).or(bw_hwm_held).unwrap_or(0.0) / 1e6,
318                        signal.map_or(0.0, |s| s.bw_avg_bps) / 1e6,
319                        signal.map_or(0, |s| s.ttfb_min.as_millis()),
320                        signal.map_or(0, |s| s.ttfb_avg.as_millis()),
321                        signal.map_or(0, |s| s.bdp_bytes()) as f64 / 1e6,
322                        stats.bytes_budget as f64 / 1e6,
323                        stats.bytes_in_use as f64 / 1e6,
324                        stats.bytes_saturation,
325                        stats.request_budget,
326                        stats.requests_in_use,
327                        stats.requests_saturation
328                    );
329                    if dropped > 0 {
330                        eprintln!(
331                            "[InFlightConcurrency] WARN: {dropped} samples dropped (queue full)"
332                        );
333                    }
334                }
335            }
336        })
337    }
338}
339
340impl Drop for ConcurrencyController {
341    fn drop(&mut self) {
342        self._control_task.abort();
343    }
344}
polars_io/cloud/concurrency/mod.rs

polars_io/cloud/concurrency/
mod.rs
