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