saluki_core/components/relays/
context.rs

1use std::sync::Arc;
2
3use saluki_common::sync::shutdown::ShutdownHandle;
4
5use crate::accounting::ComponentRegistry;
6use crate::health::Health;
7use crate::{
8    components::{ComponentContext, ComponentSpawner},
9    topology::{PayloadsDispatcher, TopologyContext},
10};
11
12struct RelayContextInner {
13    topology_context: TopologyContext,
14    component_context: ComponentContext,
15    component_registry: ComponentRegistry,
16    dispatcher: PayloadsDispatcher,
17    spawner: ComponentSpawner,
18}
19
20/// Relay context.
21pub struct RelayContext {
22    shutdown_handle: Option<ShutdownHandle>,
23    health_handle: Option<Health>,
24    inner: Arc<RelayContextInner>,
25}
26
27impl RelayContext {
28    /// Creates a new `RelayContext`.
29    pub fn new(
30        topology_context: &TopologyContext, component_context: &ComponentContext,
31        component_registry: ComponentRegistry, health_handle: Health, dispatcher: PayloadsDispatcher,
32        spawner: ComponentSpawner,
33    ) -> Self {
34        Self {
35            shutdown_handle: None,
36            health_handle: Some(health_handle),
37            inner: Arc::new(RelayContextInner {
38                topology_context: topology_context.clone(),
39                component_context: component_context.clone(),
40                component_registry,
41                dispatcher,
42                spawner,
43            }),
44        }
45    }
46
47    /// Installs the shutdown handle for this relay context.
48    ///
49    /// Called once by the runtime, before the component runs, with the shutdown signal of the
50    /// component's dedicated supervisor.
51    pub(crate) fn set_shutdown_handle(&mut self, shutdown_handle: ShutdownHandle) {
52        self.shutdown_handle = Some(shutdown_handle);
53    }
54
55    /// Installs the shutdown handle for this relay context, for tests.
56    ///
57    /// The topology runtime does this itself before a relay runs, using the shutdown signal of the component's
58    /// dedicated supervisor. A test that drives a relay through a real shutdown has to stand in for it.
59    #[cfg(any(test, feature = "test-util"))]
60    pub fn set_shutdown_handle_for_test(&mut self, shutdown_handle: ShutdownHandle) {
61        self.set_shutdown_handle(shutdown_handle);
62    }
63
64    /// Consumes the shutdown handle of this relay context.
65    ///
66    /// # Panics
67    ///
68    /// Panics if the shutdown handle has already been taken.
69    pub fn take_shutdown_handle(&mut self) -> ShutdownHandle {
70        self.shutdown_handle.take().expect("shutdown handle already taken")
71    }
72
73    /// Consumes the health handle of this relay context.
74    ///
75    /// # Panics
76    ///
77    /// Panics if the health handle has already been taken.
78    pub fn take_health_handle(&mut self) -> Health {
79        self.health_handle.take().expect("health handle already taken")
80    }
81
82    /// Returns a reference to the topology context.
83    pub fn topology_context(&self) -> &TopologyContext {
84        &self.inner.topology_context
85    }
86
87    /// Returns a reference to the component context.
88    pub fn component_context(&self) -> &ComponentContext {
89        &self.inner.component_context
90    }
91
92    /// Returns a reference to the component registry.
93    pub fn component_registry(&self) -> &ComponentRegistry {
94        &self.inner.component_registry
95    }
96
97    /// Returns a reference to the payloads dispatcher.
98    pub fn dispatcher(&self) -> &PayloadsDispatcher {
99        &self.inner.dispatcher
100    }
101
102    /// Returns a spawner for supervised child tasks belonging to this component.
103    ///
104    /// All child tasks spawned through this mechanism are tied to the lifecycle of the component itself, such that
105    /// they're automatically shutdown/stopped when the component is stopped during topology shutdown, etc.
106    pub fn spawner(&self) -> &ComponentSpawner {
107        &self.inner.spawner
108    }
109}
110
111impl Clone for RelayContext {
112    fn clone(&self) -> Self {
113        Self {
114            shutdown_handle: None,
115            health_handle: None,
116            inner: self.inner.clone(),
117        }
118    }
119}