query_runtime/runtime/
io_lifecycle.rs

1//! Revocable process-I/O lifetime used by server generations.
2
3use std::{
4    sync::{
5        Arc, Condvar, Mutex,
6        atomic::{AtomicUsize, Ordering},
7    },
8    time::Instant,
9};
10
11use query_planning::error::{QueryError, QueryErrorCode, Result};
12
13const IO_LIFECYCLE_CLOSED: usize = 1usize << (usize::BITS - 1);
14const IO_LIFECYCLE_ACTIVE_MASK: usize = IO_LIFECYCLE_CLOSED - 1;
15
16#[derive(Debug, Default)]
17struct IoLifecycleInner {
18    lifecycle: AtomicUsize,
19    wait: Mutex<()>,
20    changed: Condvar,
21}
22
23/// One revocable I/O generation shared by native execution contexts.
24///
25/// Acquisition and close update one atomic closed-bit/count word, so an
26/// operation is either counted before close or rejected after close. Cloning
27/// this handle does not itself count as active I/O.
28#[derive(Clone, Debug, Default)]
29pub struct IoLifecycle {
30    inner: Arc<IoLifecycleInner>,
31}
32
33/// Counted permission to execute one file-backed native I/O operation.
34#[derive(Debug)]
35pub struct IoLifecycleLease {
36    lifecycle: IoLifecycle,
37}
38
39impl IoLifecycle {
40    pub fn new() -> Self {
41        Self::default()
42    }
43
44    pub fn acquire(&self, operation: &'static str) -> Result<IoLifecycleLease> {
45        let mut lifecycle = self.inner.lifecycle.load(Ordering::Acquire);
46        loop {
47            if lifecycle & IO_LIFECYCLE_CLOSED != 0 {
48                return Err(io_lifecycle_closed(operation));
49            }
50            if lifecycle & IO_LIFECYCLE_ACTIVE_MASK == IO_LIFECYCLE_ACTIVE_MASK {
51                return Err(QueryError::new(
52                    QueryErrorCode::Internal,
53                    operation,
54                    "server I/O generation lease count overflow",
55                )
56                .with_fact("reason", "io_lifecycle_lease_count_overflow"));
57            }
58            match self.inner.lifecycle.compare_exchange_weak(
59                lifecycle,
60                lifecycle + 1,
61                Ordering::AcqRel,
62                Ordering::Acquire,
63            ) {
64                Ok(_) => return Ok(IoLifecycleLease { lifecycle: self.clone() }),
65                Err(observed) => lifecycle = observed,
66            }
67        }
68    }
69
70    /// Close this generation without blocking. Repeated close calls are safe.
71    pub fn close(&self) {
72        self.inner.lifecycle.fetch_or(IO_LIFECYCLE_CLOSED, Ordering::AcqRel);
73        self.inner.changed.notify_all();
74    }
75
76    pub(crate) fn active_leases(&self) -> usize {
77        self.inner.lifecycle.load(Ordering::Acquire) & IO_LIFECYCLE_ACTIVE_MASK
78    }
79
80    /// Close and wait until every lease acquired before close has dropped.
81    pub fn close_and_wait_for_idle_until(&self, deadline: Instant, operation: &'static str) -> Result<()> {
82        self.close();
83        let mut wait = match super::deadline::mutex_lock_until(&self.inner.wait, deadline) {
84            Ok(wait) => wait,
85            Err(super::deadline::DeadlineMutexLockError::Poisoned) => return Err(io_lifecycle_lock_error(operation)),
86            Err(super::deadline::DeadlineMutexLockError::TimedOut) => {
87                return Err(io_lifecycle_timeout(operation, None));
88            },
89        };
90        let mut waited_for_idle = false;
91        loop {
92            let active_leases = self.active_leases();
93            if active_leases == 0 {
94                if waited_for_idle && super::deadline::deadline_expired(deadline) {
95                    return Err(io_lifecycle_timeout(operation, Some(0)));
96                }
97                return Ok(());
98            }
99            if super::deadline::deadline_expired(deadline) {
100                return Err(io_lifecycle_timeout(operation, Some(active_leases)));
101            }
102            let Some(remaining) = super::deadline::deadline_remaining(deadline) else {
103                return Err(io_lifecycle_timeout(operation, Some(active_leases)));
104            };
105            waited_for_idle = true;
106            let (next, _timeout) = self.inner.changed.wait_timeout(wait, remaining).map_err(|_error| {
107                QueryError::new(
108                    QueryErrorCode::Internal,
109                    operation,
110                    "native I/O lifecycle condition variable was poisoned during shutdown",
111                )
112                .with_fact("reason", "io_lifecycle_wait_poisoned")
113            })?;
114            wait = next;
115        }
116    }
117
118    fn release(&self) {
119        let _wait = self.inner.wait.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
120        let previous = self.inner.lifecycle.fetch_sub(1, Ordering::AcqRel);
121        debug_assert!(previous & IO_LIFECYCLE_ACTIVE_MASK > 0, "native I/O lifecycle lease count underflow");
122        let lifecycle = previous - 1;
123        if lifecycle & IO_LIFECYCLE_CLOSED != 0 && lifecycle & IO_LIFECYCLE_ACTIVE_MASK == 0 {
124            self.inner.changed.notify_all();
125        }
126    }
127}
128
129impl Drop for IoLifecycleLease {
130    fn drop(&mut self) {
131        self.lifecycle.release();
132    }
133}
134
135fn io_lifecycle_closed(operation: &'static str) -> QueryError {
136    QueryError::new(QueryErrorCode::Cancelled, operation, "server I/O generation is closed")
137        .with_fact("reason", "io_lifecycle_closed")
138}
139
140fn io_lifecycle_lock_error(operation: &'static str) -> QueryError {
141    QueryError::new(QueryErrorCode::Internal, operation, "native I/O lifecycle mutex was poisoned")
142        .with_fact("reason", "io_lifecycle_mutex_poisoned")
143}
144
145fn io_lifecycle_timeout(operation: &'static str, active_leases: Option<usize>) -> QueryError {
146    let error = QueryError::new(
147        QueryErrorCode::ServiceOverloaded,
148        operation,
149        "timed out waiting for server I/O generation to become idle",
150    )
151    .with_fact("reason", "io_lifecycle_shutdown_timeout");
152    match active_leases {
153        Some(active_leases) => error.with_fact("active_io_leases", active_leases),
154        None => error,
155    }
156}
157
158#[cfg(test)]
159mod tests {
160    use std::time::{Duration, Instant};
161
162    use super::*;
163
164    #[test]
165    fn close_rejects_new_io_and_waits_for_prior_lease() {
166        let lifecycle = IoLifecycle::new();
167        let lease = lifecycle.acquire("test.io.start").unwrap();
168        lifecycle.close();
169        let error = lifecycle.acquire("test.io.after_close").unwrap_err();
170        assert_eq!(error.code(), QueryErrorCode::Cancelled);
171        assert_eq!(error.fact_value("reason"), Some("io_lifecycle_closed"));
172        let timeout = lifecycle.close_and_wait_for_idle_until(Instant::now(), "test.io.wait").unwrap_err();
173        assert_eq!(timeout.fact_value("reason"), Some("io_lifecycle_shutdown_timeout"));
174        assert_eq!(timeout.fact_value("active_io_leases"), Some("1"));
175        drop(lease);
176        lifecycle.close_and_wait_for_idle_until(Instant::now() + Duration::from_secs(1), "test.io.wait").unwrap();
177    }
178
179    #[test]
180    fn contended_state_lock_cannot_turn_post_deadline_idle_into_success() {
181        let lifecycle = IoLifecycle::new();
182        let state = lifecycle.inner.wait.lock().unwrap();
183        let deadline = Instant::now() + Duration::from_millis(20);
184        let waiting_lifecycle = lifecycle.clone();
185        let (result_tx, result_rx) = std::sync::mpsc::channel();
186        let wait = std::thread::spawn(move || {
187            result_tx
188                .send(waiting_lifecycle.close_and_wait_for_idle_until(deadline, "test.io.contended_shutdown"))
189                .unwrap();
190        });
191
192        let result = result_rx.recv_timeout(Duration::from_millis(200));
193        drop(state);
194        wait.join().unwrap();
195
196        let error = result
197            .expect("I/O shutdown must return its mutex timeout before the contended lock is released")
198            .expect_err("idle observed after the deadline must not be accepted");
199        assert_eq!(error.fact_value("reason"), Some("io_lifecycle_shutdown_timeout"));
200        assert_eq!(
201            lifecycle.acquire("test.io.after_mutex_timeout").unwrap_err().fact_value("reason"),
202            Some("io_lifecycle_closed"),
203            "mutex timeout must leave Parquet I/O fail-closed"
204        );
205    }
206
207    #[test]
208    fn close_signal_is_nonblocking_and_fail_closed_during_state_lock_contention() {
209        let lifecycle = IoLifecycle::new();
210        let state = lifecycle.inner.wait.lock().unwrap();
211        let closing_lifecycle = lifecycle.clone();
212        let (closed_tx, closed_rx) = std::sync::mpsc::channel();
213        let close = std::thread::spawn(move || {
214            closing_lifecycle.close();
215            closed_tx.send(()).unwrap();
216        });
217
218        closed_rx.recv_timeout(Duration::from_millis(50)).expect("close signal must not wait for the I/O state mutex");
219        drop(state);
220        close.join().unwrap();
221        assert_eq!(
222            lifecycle.acquire("test.io.after_nonblocking_close").unwrap_err().fact_value("reason"),
223            Some("io_lifecycle_closed")
224        );
225    }
226}
query_runtime/runtime/io_lifecycle.rs

query_runtime/runtime/
io_lifecycle.rs
