polars_io/cloud/concurrency/

mod.rs

//! Adaptive in-flight concurrency controller for cloud IO (Input only).
//!
//! Admission control for concurrency uses two budgets:
//! - A (primary) bytes-based budget to model the bandwidth-delay product (BDP)
//! - A (secondary) count-based budget to limit the number of in-flight requests
//!
//! The bytes-based budget models the BDP as
//!   BDP = BW_max * TTFB_est
//!
//! Three components cooperate:
//! - Model: records IO observations and models the network (BW_max, TTFB_est, BDP)
//! - Regime: state machine driving the admission (Init / RampUp / Stable / ProbeUp)
//! - Admission: admission control, enforces byte-based + request-based budgets

// Loosely based on BBR: Congestion-Based Congestion Control
// see https://queue.acm.org/detail.cfm?id=3022184

mod admission;
mod model;
mod regime;

use std::num::NonZeroU64;
use std::sync::Arc;
use std::sync::atomic::AtomicU64;
use std::sync::atomic::Ordering::Relaxed;
use std::time::{Duration, Instant};

pub use admission::{InFlightBudget, InFlightPermit, InFlightStats};
use crossbeam_queue::ArrayQueue;
pub use model::Model;
use polars_core::runtime::ASYNC;
use polars_utils::relaxed_cell::RelaxedCell;
pub use regime::{Regime, RegimeState};

// Number of samples in the queue, which gets drained on every tick.
// At 50k requests per second and 100 ms tick window, we need 5k.
const SAMPLE_QUEUE_CAPACITY: usize = 8192;

use crate::cloud::concurrency_config::get_random_access_chunk_size;
use crate::cloud::http_rate_limit::PacingBudget;

#[derive(Clone, Copy, Debug)]
pub struct IoSample {
    pub n_bytes: u64,
    // Time-to-first-byte.
    pub ttfb: Duration,
    // TODO: Factor out as we only care about per-tick_window stats.
    pub completion_time: Instant,
}

#[derive(Debug, Clone)]
pub struct ControllerConfig {
    // Sliding window over which the most recent round-trip-time (RTT) and bandwidth (BW)
    // will be calculated. Also acts as the retention window.
    window: Duration,
    // Byte-based budget during the Init phase, and as the base for the RampUp phase.
    init_byte_budget: u64,
    // Lower limit for the byte-based budget - needed to avoid deadlock.
    floor_byte_budget: u64,
    // Count-based request budget.
    request_budget: u64,
    // Lower limit for the count-based budget.
    floor_request_budget: u64,
    // Controller regime update frequency.
    control_interval: Duration,
    // Total budget only resizes if the relative changes exceeds this threshold
    budget_resize_threshold: f64,
}

impl Default for ControllerConfig {
    fn default() -> Self {
        // Only used for bytes-based budget.
        let target_chunk_size = get_random_access_chunk_size() as u64;
        Self {
            window: Duration::from_millis(1000),

            // Byte-based budget during the ramp-up phase.
            // Starting too low results in lost opportunity (time) during ramp-up.
            // Starting too high results in early congestion, delayed completion of the first chunk,
            // and inflated bandwidth estimation.
            //
            // Some BDP numbers for reference:
            //   1 Gbps x 20 ms = 2.5 MB
            //   1 Gbps x 50 ms = 6.25 MB
            //   10 Gbps x 50 ms = 62.5 MB
            //   100 Gbps x 50 ms = 625 MB
            init_byte_budget: get_init_byte_budget(target_chunk_size),

            // Byte-based budget floor.
            // Must be >=larger than target_chunk_size to avoid potential deadlock.
            floor_byte_budget: target_chunk_size,

            // Count-based budget.
            request_budget: get_request_budget(),
            floor_request_budget: get_floor_request_budget(),
            control_interval: Duration::from_millis(100),
            budget_resize_threshold: 0.05,
        }
    }
}

/// Max number of bytes concurrently in flight during the init and start of rampup phase.
fn get_init_byte_budget(target_chunk_size: u64) -> u64 {
    let init_byte_budget = std::env::var("POLARS_INFLIGHT_INIT_BYTE_BUDGET")
        .map(|x| {
            x.parse::<NonZeroU64>()
                .unwrap_or_else(|_| {
                    panic!("invalid value for POLARS_INFLIGHT_INIT_BYTE_BUDGET: {x}")
                })
                .get()
        })
        .unwrap_or_else(|_| {
            // This should be lower than the expected BDP so it can ramp-up, but
            // too low a value delays the transition to stable.
            // Heuristic: higher bandwidth is expected on larger instances.
            let n = polars_config::config().max_threads() as u64;
            n.div_ceil(8).max(4) * target_chunk_size
        })
        .max(1);

    if init_byte_budget < target_chunk_size {
        panic!("in-flight byte budget init must be larger than the target_chunk_size");
    }

    init_byte_budget
}

/// Maximum number of requests concurrently in flight.
pub fn get_request_budget() -> u64 {
    // Since object_store/reqwest use HTTP/1 with a connection pool, this value controls the
    // max concurrent TCP sessions to S3 for the pipeline.
    // When modifying this value, consider the max_thread count configuration(s), the OS limitations
    // (e.g., ulimit -n), and any cloud infrastructure limitations.
    std::env::var("POLARS_INFLIGHT_REQUEST_BUDGET")
        .map(|x| {
            x.parse::<NonZeroU64>()
                .unwrap_or_else(|_| panic!("invalid value for POLARS_INFLIGHT_REQUEST_BUDGET: {x}"))
                .get()
        })
        .unwrap_or(512)
        .max(1)
}

/// Minimum number of requests concurrently in flight, if demand is there.
pub fn get_floor_request_budget() -> u64 {
    // Since object_store/reqwest use HTTP/1 with a connection pool, this value controls the
    // max concurrent TCP sessions to S3 for the pipeline.
    // When modifying this value, consider the max_thread count configuration(s), the OS limitations
    // (e.g., ulimit -n), and any cloud infrastructure limitations.
    std::env::var("POLARS_INFLIGHT_FLOOR_REQUEST_BUDGET")
        .map(|x| {
            x.parse::<NonZeroU64>()
                .unwrap_or_else(|_| {
                    panic!("invalid value for POLARS_INFLIGHT_FLOOR_REQUEST_BUDGET: {x}")
                })
                .get()
        })
        .unwrap_or(polars_config::config().max_threads() as u64)
        .max(1)
}

/// Windowed round-trip time of 0-byte metadata (HEAD) requests.
/// Diagnostic only, not to be used for BDP estimation.
#[derive(Debug)]
pub struct HeadRttChannel {
    min_ns: AtomicU64,
    sum_ns: AtomicU64,
    count: AtomicU64,
    /// Never reset, unlike `count`: distinguishes the first request of the process from
    /// the first of a window.
    total: AtomicU64,
}

/// One control tick's worth of [`HeadRttChannel`] observations.
#[derive(Clone, Copy, Debug)]
pub struct HeadRttWindow {
    pub min: Option<Duration>,
    pub avg: Option<Duration>,
    pub count: u64,
}

impl HeadRttChannel {
    fn new() -> Self {
        Self {
            min_ns: AtomicU64::new(u64::MAX),
            sum_ns: AtomicU64::new(0),
            count: AtomicU64::new(0),
            total: AtomicU64::new(0),
        }
    }

    /// Hot path: four relaxed ops, no queue.
    fn record(&self, rtt: Duration) {
        let ns = rtt.as_nanos() as u64;
        self.min_ns.fetch_min(ns, Relaxed);
        self.sum_ns.fetch_add(ns, Relaxed);
        self.count.fetch_add(1, Relaxed);

        if self.total.fetch_add(1, Relaxed) == 0 && polars_config::config().verbose() {
            eprintln!(
                "[InFlightConcurrency]: observed first RTT sample (metadata): {} ms",
                rtt.as_millis()
            );
        }
    }

    /// Read and reset.
    fn take(&self) -> HeadRttWindow {
        let min_ns = self.min_ns.swap(u64::MAX, Relaxed);
        let sum_ns = self.sum_ns.swap(0, Relaxed);
        let count = self.count.swap(0, Relaxed);

        HeadRttWindow {
            min: (min_ns != u64::MAX).then(|| Duration::from_nanos(min_ns)),
            avg: (count > 0).then(|| Duration::from_nanos(sum_ns / count)),
            count,
        }
    }
}

#[derive(Debug)]
pub struct ConcurrencyController {
    config: ControllerConfig,
    sample_queue: Arc<ArrayQueue<IoSample>>,
    samples_dropped: Arc<RelaxedCell<u64>>,
    head_rtt: Arc<HeadRttChannel>,
    inflight_budget: Arc<InFlightBudget>,
    _control_task: tokio::task::JoinHandle<()>,
}

impl ConcurrencyController {
    pub fn new(config: ControllerConfig, pacing_budget: Option<PacingBudget>) -> Self {
        let sample_queue = Arc::new(ArrayQueue::new(SAMPLE_QUEUE_CAPACITY));
        let samples_dropped = Arc::new(RelaxedCell::new_u64(0));
        let head_rtt = Arc::new(HeadRttChannel::new());

        let inflight_budget = Arc::new(InFlightBudget::new(
            config.init_byte_budget,
            config.floor_byte_budget,
            config.request_budget,
            config.floor_request_budget,
        ));

        let control_task = Self::spawn_control_loop(
            sample_queue.clone(),
            samples_dropped.clone(),
            head_rtt.clone(),
            inflight_budget.clone(),
            config.clone(),
            pacing_budget,
        );

        Self {
            config,
            sample_queue,
            samples_dropped,
            head_rtt,
            inflight_budget,
            _control_task: control_task,
        }
    }

    pub fn config(&self) -> &ControllerConfig {
        &self.config
    }

    /// Record IO for a completed data request. Hot path.
    pub fn record_io(&self, sample: IoSample) {
        if self.sample_queue.push(sample).is_err() {
            // Queue full: drop. Samples are statistics is considered acceptable.
            self.samples_dropped.fetch_add(1);
        }
    }

    /// Record the round-trip time of a completed 0-byte metadata request. Hot path.
    pub fn record_head_rtt(&self, rtt: Duration) {
        self.head_rtt.record(rtt);
    }

    pub fn inflight_budget(&self) -> &Arc<InFlightBudget> {
        &self.inflight_budget
    }

    pub async fn acquire(&self, bytes: u64) -> InFlightPermit {
        self.inflight_budget.acquire(bytes).await
    }

    fn spawn_control_loop(
        sample_queue: Arc<ArrayQueue<IoSample>>,
        samples_dropped: Arc<RelaxedCell<u64>>,
        head_rtt: Arc<HeadRttChannel>,
        admission: Arc<InFlightBudget>,
        config: ControllerConfig,
        pacing_budget: Option<PacingBudget>,
    ) -> tokio::task::JoinHandle<()> {
        if polars_config::config().verbose() {
            eprintln!(
                "[InFlightConcurrency]: spawn control loop: control_interval: {}ms",
                config.control_interval.as_millis()
            );
        }
        ASYNC.spawn(async move {
            let mut model = Model::new(config.window);
            let mut regime = Regime::new(Instant::now());

            let mut ticker = tokio::time::interval(config.control_interval);
            ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);

            loop {
                ticker.tick().await;
                let now = Instant::now();

                // Limit concurrency to the pacing budget from the rate-limiter.
                if let Some(ref pacing_budget) = pacing_budget {
                    let rate = pacing_budget.rate();
                    let horizon_s = pacing_budget.horizon().as_secs_f64();
                    let request_budget = rate * horizon_s;
                    admission.resize_request_budget(request_budget as u64);
                }

                // Update model statistics and step regime.
                let (state, signal, dropped, bw_hwm_held) = {
                    for _ in 0..SAMPLE_QUEUE_CAPACITY {
                        let Some(s) = sample_queue.pop() else { break };
                        model.record(s);
                    }
                    let dropped = samples_dropped.swap(0);
                    model.update(now);
                    let signal = model.signal();
                    let state = regime.step(signal, now);
                    let bw_hwm_held = model.bw_hwm_bps();
                    (state, signal, dropped, bw_hwm_held)
                };

                if !matches!(state, RegimeState::WarmIdle { .. }) {
                    // Compute base BDP
                    let base_byte_budget = match (state, signal) {
                        (RegimeState::Init, _) | (_, None) => config.init_byte_budget,
                        (_, Some(signal)) => signal.bdp_bytes().max(config.init_byte_budget),
                    };

                    // Compute target BDP using the gain multiplier. This is similar to BBR cwnd_gain.
                    let gain = match state {
                        RegimeState::Init => 1.0,
                        RegimeState::RampUp { .. } => 2.0,
                        // NOTE: >> 1.0 for the purpose of absorbing environment noise.
                        RegimeState::Stable => 2.0,
                        RegimeState::ProbeUp { .. } => 3.0,
                        // Unreachable.
                        RegimeState::WarmIdle { .. } => 1.0,
                    };
                    let target_byte_budget = (base_byte_budget as f64 * gain) as u64;

                    // Resize if needed.
                    let current_byte_budget = admission.current_byte_budget();
                    let threshold = config.budget_resize_threshold;
                    let should_resize = match current_byte_budget {
                        0 => target_byte_budget > 0,
                        current => {
                            let ratio = target_byte_budget as f64 / current as f64;
                            ratio < (1.0 - threshold) || ratio > (1.0 + threshold)
                        },
                    };

                    if should_resize {
                        admission.resize_byte_budget(target_byte_budget);
                    }
                }

                // Drained every tick. Diagnostic only.
                let head_rtt_window = head_rtt.take();

                // Log snapshot.
                if std::env::var("POLARS_LOG_CONCURRENCY").is_ok() {
                    let stats = admission.stats();
                    eprintln!(
                        "[InFlightConcurrency {}] regime={}, \
                        bw_hwm={:.1} MB/s, \
                        bw_avg={:.1} MB/s, \
                        rtt_min={:.1} ms, \
                        rtt_avg={:.1} ms, \
                        head_rtt_min={:.1} ms, \
                        head_rtt_avg={:.1} ms, \
                        head_n={}, \
                        bdp_obs={:.1} MB, \
                        bytes_budget={:.1} MB, \
                        bytes_in_use={:.1} MB, \
                        bytes_sat={:.2}, \
                        req_budget={}, \
                        req_in_use={}, \
                        req_sat={:.2}",
                        chrono::Utc::now(),
                        state.label(),
                        signal.map(|s| s.bw_hwm_bps).or(bw_hwm_held).unwrap_or(0.0) / 1e6,
                        signal.map_or(0.0, |s| s.bw_avg_bps) / 1e6,
                        signal.map_or(0, |s| s.ttfb_min.as_millis()),
                        signal.map_or(0, |s| s.ttfb_avg.as_millis()),
                        head_rtt_window.min.map_or(0.0, |d| d.as_secs_f64() * 1e3),
                        head_rtt_window.avg.map_or(0.0, |d| d.as_secs_f64() * 1e3),
                        head_rtt_window.count,
                        signal.map_or(0, |s| s.bdp_bytes()) as f64 / 1e6,
                        stats.bytes_budget as f64 / 1e6,
                        stats.bytes_in_use as f64 / 1e6,
                        stats.bytes_saturation,
                        stats.request_budget,
                        stats.requests_in_use,
                        stats.requests_saturation
                    );
                    if dropped > 0 {
                        eprintln!(
                            "[InFlightConcurrency] WARN: {dropped} samples dropped (queue full)"
                        );
                    }
                }
            }
        })
    }
}

impl Drop for ConcurrencyController {
    fn drop(&mut self) {
        self._control_task.abort();
    }
}