iceberg/runtime/
mod.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18// This module contains the async runtime abstraction for iceberg.
19
20use std::fmt;
21use std::future::Future;
22use std::pin::Pin;
23use std::task::{Context, Poll};
24
25use tokio::task::JoinHandle as TokioJoinHandle;
26
27use crate::{Error, ErrorKind, Result};
28
29/// Wrapper around tokio's [`TokioJoinHandle`] that converts task failures into
30/// [`crate::Error`].
31///
32/// Tokio's `JoinHandle<T>` resolves to `Result<T, JoinError>`, where a
33/// [`tokio::task::JoinError`] means the task either panicked or was cancelled
34/// (typically from runtime shutdown or `abort`).
35/// Both are surfaced here as `ErrorKind::Unexpected` with the original
36/// [`tokio::task::JoinError`] preserved as the source.
37pub struct JoinHandle<T>(TokioJoinHandle<T>);
38
39impl<T> JoinHandle<T> {
40    /// Ties the task's lifetime to this handle: see [`AbortOnDropHandle`].
41    pub(crate) fn abort_on_drop(self) -> AbortOnDropHandle<T> {
42        AbortOnDropHandle(self)
43    }
44}
45
46impl<T> Unpin for JoinHandle<T> {}
47
48impl<T: Send + 'static> Future for JoinHandle<T> {
49    type Output = Result<T>;
50
51    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
52        Pin::new(&mut self.get_mut().0).poll(cx).map(|r| {
53            r.map_err(|e| Error::new(ErrorKind::Unexpected, "spawned task failed").with_source(e))
54        })
55    }
56}
57
58/// A [`JoinHandle`] that aborts its task when dropped.
59///
60/// Dropping a plain `JoinHandle` leaves the task running with nobody waiting for
61/// its result. Use this instead when only the caller wants that result, so that
62/// dropping the caller stops the work.
63pub(crate) struct AbortOnDropHandle<T>(JoinHandle<T>);
64
65impl<T> Drop for AbortOnDropHandle<T> {
66    fn drop(&mut self) {
67        self.0.0.abort();
68    }
69}
70
71impl<T: Send + 'static> Future for AbortOnDropHandle<T> {
72    type Output = Result<T>;
73
74    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
75        Pin::new(&mut self.get_mut().0).poll(cx)
76    }
77}
78
79/// Handle to a single tokio runtime.
80///
81/// Wraps a [`tokio::runtime::Handle`], which is cheap to clone. The caller is
82/// responsible for keeping the underlying runtime alive while this handle is
83/// in use; spawning on a shut-down runtime will surface as a [`tokio::task::JoinError`]
84/// via the returned [`JoinHandle`].
85#[derive(Clone)]
86pub struct RuntimeHandle {
87    handle: tokio::runtime::Handle,
88}
89
90impl fmt::Debug for RuntimeHandle {
91    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
92        f.debug_struct("RuntimeHandle").finish()
93    }
94}
95
96impl RuntimeHandle {
97    fn from_tokio_handle(handle: tokio::runtime::Handle) -> Self {
98        Self { handle }
99    }
100
101    /// Spawn an async task.
102    pub fn spawn<F>(&self, future: F) -> JoinHandle<F::Output>
103    where
104        F: Future + Send + 'static,
105        F::Output: Send + 'static,
106    {
107        JoinHandle(self.handle.spawn(future))
108    }
109
110    /// Spawn a blocking task.
111    pub fn spawn_blocking<F, T>(&self, f: F) -> JoinHandle<T>
112    where
113        F: FnOnce() -> T + Send + 'static,
114        T: Send + 'static,
115    {
116        JoinHandle(self.handle.spawn_blocking(f))
117    }
118}
119
120/// Iceberg's runtime abstraction.
121///
122/// Contains separate handles for IO-bound and CPU-bound work. When constructed
123/// with a single tokio runtime, both `io()` and `cpu()` route to the same one.
124/// Use [`Runtime::new_with_split`] to provide dedicated runtimes for each
125/// category.
126///
127/// # Lifetime
128///
129/// A `Runtime` stores only `tokio::runtime::Handle`s (weak references).
130/// The caller owns the tokio runtime's lifetime. If the underlying runtime is
131/// dropped while iceberg is still using it, subsequent spawns will surface as
132/// task cancellation errors via the returned `JoinHandle`.
133///
134/// Cloning is cheap.
135#[derive(Clone)]
136pub struct Runtime {
137    io: RuntimeHandle,
138    cpu: RuntimeHandle,
139}
140
141impl fmt::Debug for Runtime {
142    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
143        f.debug_struct("Runtime").finish()
144    }
145}
146
147impl Runtime {
148    /// Create a Runtime backed by a single tokio runtime for all work.
149    pub fn new(runtime: &tokio::runtime::Runtime) -> Self {
150        let handle = RuntimeHandle::from_tokio_handle(runtime.handle().clone());
151        Self {
152            io: handle.clone(),
153            cpu: handle,
154        }
155    }
156
157    /// Create a Runtime with separate tokio runtimes for IO and CPU work.
158    pub fn new_with_split(
159        io_runtime: &tokio::runtime::Runtime,
160        cpu_runtime: &tokio::runtime::Runtime,
161    ) -> Self {
162        Self {
163            io: RuntimeHandle::from_tokio_handle(io_runtime.handle().clone()),
164            cpu: RuntimeHandle::from_tokio_handle(cpu_runtime.handle().clone()),
165        }
166    }
167
168    /// Borrows the tokio runtime the caller is currently running in.
169    ///
170    /// Panics if called outside a tokio runtime context. Use
171    /// [`Runtime::try_current`] for a fallible version.
172    ///
173    /// Iceberg never implicitly spawns its own runtime; callers outside a
174    /// tokio context must construct one explicitly via [`Runtime::new`] or
175    /// [`Runtime::new_with_split`].
176    pub fn current() -> Self {
177        Self::try_current().expect(
178            "Runtime::current() called outside a tokio runtime context. \
179             Call it from within #[tokio::main] / #[tokio::test], or construct \
180             a Runtime explicitly via Runtime::new / Runtime::new_with_split.",
181        )
182    }
183
184    /// Fallible variant of [`Runtime::current`]. Returns an error if no tokio
185    /// runtime is available in the current context.
186    pub fn try_current() -> Result<Self> {
187        let handle = tokio::runtime::Handle::try_current().map_err(|e| {
188            Error::new(
189                ErrorKind::Unexpected,
190                "no tokio runtime in context; call Runtime::try_current() \
191                 from within a tokio runtime, or construct a Runtime explicitly \
192                 via Runtime::new / Runtime::new_with_split",
193            )
194            .with_source(e)
195        })?;
196        let rh = RuntimeHandle::from_tokio_handle(handle);
197        Ok(Self {
198            io: rh.clone(),
199            cpu: rh,
200        })
201    }
202
203    /// Handle for IO-bound work (network fetches, file reads).
204    pub fn io(&self) -> &RuntimeHandle {
205        &self.io
206    }
207
208    /// Handle for CPU-bound work (decoding, predicate eval, projection).
209    pub fn cpu(&self) -> &RuntimeHandle {
210        &self.cpu
211    }
212}
213
214#[cfg(test)]
215mod tests {
216    use super::*;
217
218    /// A test harness that owns a tokio runtime and exposes a `Runtime` handle
219    /// plus a `block_on` helper for sync test bodies.
220    struct TestRuntime {
221        tokio: tokio::runtime::Runtime,
222        rt: Runtime,
223    }
224
225    impl TestRuntime {
226        fn new() -> Self {
227            let tokio = tokio::runtime::Builder::new_multi_thread()
228                .enable_all()
229                .build()
230                .expect("Failed to build tokio runtime");
231            let rt = Runtime::new(&tokio);
232            Self { tokio, rt }
233        }
234
235        fn block_on<F: Future>(&self, f: F) -> F::Output {
236            self.tokio.block_on(f)
237        }
238    }
239
240    #[test]
241    fn test_runtime_spawn_io() {
242        let h = TestRuntime::new();
243        let handle = h.rt.io().spawn(async { 1 + 1 });
244        assert_eq!(h.block_on(handle).unwrap(), 2);
245    }
246
247    #[test]
248    fn test_runtime_spawn_cpu() {
249        let h = TestRuntime::new();
250        let handle = h.rt.cpu().spawn(async { 3 + 4 });
251        assert_eq!(h.block_on(handle).unwrap(), 7);
252    }
253
254    #[test]
255    fn test_runtime_spawn_blocking() {
256        let h = TestRuntime::new();
257        let handle = h.rt.cpu().spawn_blocking(|| 1 + 1);
258        assert_eq!(h.block_on(handle).unwrap(), 2);
259    }
260
261    #[test]
262    fn test_runtime_new_with_custom_runtime() {
263        let h = TestRuntime::new();
264        let handle = h.rt.io().spawn(async { 42 });
265        assert_eq!(h.block_on(handle).unwrap(), 42);
266    }
267
268    #[test]
269    fn test_runtime_split_uses_separate_handles() {
270        let io_rt = tokio::runtime::Builder::new_multi_thread()
271            .enable_all()
272            .build()
273            .unwrap();
274        let cpu_rt = tokio::runtime::Builder::new_multi_thread()
275            .enable_all()
276            .build()
277            .unwrap();
278        let rt = Runtime::new_with_split(&io_rt, &cpu_rt);
279        // Spawn on each and confirm both are distinct live runtimes. We use
280        // `io_rt`/`cpu_rt` directly to `block_on` since our `Runtime` doesn't
281        // expose one.
282        let io_result = io_rt.block_on(async { rt.io().spawn(async { "io" }).await.unwrap() });
283        let cpu_result = cpu_rt.block_on(async { rt.cpu().spawn(async { "cpu" }).await.unwrap() });
284        assert_eq!(io_result, "io");
285        assert_eq!(cpu_result, "cpu");
286    }
287
288    #[test]
289    fn test_runtime_clone() {
290        let h = TestRuntime::new();
291        let rt2 = h.rt.clone();
292        let handle = rt2.io().spawn(async { 5 });
293        assert_eq!(h.block_on(handle).unwrap(), 5);
294    }
295
296    #[test]
297    fn test_runtime_debug() {
298        let h = TestRuntime::new();
299        let debug_str = format!("{:?}", h.rt);
300        assert!(debug_str.contains("Runtime"));
301    }
302
303    #[tokio::test(flavor = "multi_thread")]
304    async fn test_try_current_in_runtime() {
305        let rt = Runtime::try_current().expect("should find current runtime");
306        let result = rt.io().spawn(async { 7 }).await.unwrap();
307        assert_eq!(result, 7);
308    }
309
310    #[test]
311    fn test_try_current_outside_runtime() {
312        let err = Runtime::try_current().expect_err("must fail outside runtime");
313        assert_eq!(err.kind(), ErrorKind::Unexpected);
314    }
315
316    /// Verifies that when the caller drops the underlying tokio runtime, a
317    /// subsequent spawn surfaces as a `JoinError` via our `JoinHandle` rather
318    /// than hanging or misbehaving.
319    #[test]
320    fn test_spawn_after_runtime_drop_errors() {
321        let driver = tokio::runtime::Builder::new_current_thread()
322            .enable_all()
323            .build()
324            .unwrap();
325        let owned = tokio::runtime::Builder::new_multi_thread()
326            .enable_all()
327            .build()
328            .unwrap();
329        let rt = Runtime::new(&owned);
330        // Drop the caller's owned runtime. Our `Runtime` only holds a `Handle`,
331        // so this shuts down the underlying tokio runtime.
332        drop(owned);
333
334        // Spawning after shutdown returns a handle that resolves to a cancelled
335        // JoinError, which our wrapper maps to iceberg::Error.
336        let handle = rt.io().spawn(async { 1 });
337        let result = driver.block_on(handle);
338        assert!(result.is_err(), "expected error after runtime shutdown");
339    }
340}
iceberg/runtime/mod.rs

iceberg/runtime/
mod.rs
