Skip to main content

vortex_io/
session.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::any::Any;
5use std::fmt::Debug;
6
7use vortex_error::VortexExpect;
8use vortex_session::SessionExt;
9use vortex_session::SessionVar;
10
11use crate::runtime::Handle;
12
13/// Session state for Vortex async runtimes.
14pub struct RuntimeSession {
15    handle: Option<Handle>,
16}
17
18impl SessionVar for RuntimeSession {
19    fn as_any(&self) -> &dyn Any {
20        self
21    }
22
23    fn as_any_mut(&mut self) -> &mut dyn Any {
24        self
25    }
26}
27
28impl Default for RuntimeSession {
29    fn default() -> Self {
30        Self {
31            handle: Handle::find(),
32        }
33    }
34}
35
36impl Debug for RuntimeSession {
37    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
38        f.debug_struct("RuntimeSession").finish_non_exhaustive()
39    }
40}
41
42/// Extension trait for accessing runtime session data.
43pub trait RuntimeSessionExt: SessionExt {
44    /// Returns a handle for this session's runtime.
45    fn handle(&self) -> Handle {
46        self.get::<RuntimeSession>().handle
47                .as_ref()
48                .vortex_expect("Runtime handle not configured in Vortex session. Please setup a `CurrentThreadRuntime`, or configure the session for `with_tokio`.")
49                .clone()
50    }
51
52    /// Configure the runtime session to use the application's Tokio runtime.
53    ///
54    /// For example, if the application is launched using `#[tokio::main]`.
55    #[cfg(feature = "tokio")]
56    fn with_tokio(self) -> Self {
57        self.get_mut::<RuntimeSession>().handle =
58            Some(crate::runtime::tokio::TokioRuntime::current());
59        self
60    }
61
62    /// Configure the runtime session to use a specific Vortex runtime handle.
63    fn with_handle(self, handle: Handle) -> Self {
64        self.get_mut::<RuntimeSession>().handle = Some(handle);
65        self
66    }
67}
68impl<S: SessionExt> RuntimeSessionExt for S {}