saluki_core/components/sources/
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,
9    topology::{EventsDispatcher, TopologyContext},
10};
11
12struct SourceContextInner {
13    topology_context: TopologyContext,
14    component_context: ComponentContext,
15    component_registry: ComponentRegistry,
16    dispatcher: EventsDispatcher,
17}
18
19/// Source context.
20pub struct SourceContext {
21    shutdown_handle: Option<ShutdownHandle>,
22    health_handle: Option<Health>,
23    inner: Arc<SourceContextInner>,
24}
25
26impl SourceContext {
27    /// Creates a new `SourceContext`.
28    pub fn new(
29        topology_context: &TopologyContext, component_context: &ComponentContext,
30        component_registry: ComponentRegistry, health_handle: Health, dispatcher: EventsDispatcher,
31    ) -> Self {
32        Self {
33            shutdown_handle: None,
34            health_handle: Some(health_handle),
35            inner: Arc::new(SourceContextInner {
36                topology_context: topology_context.clone(),
37                component_context: component_context.clone(),
38                component_registry,
39                dispatcher,
40            }),
41        }
42    }
43
44    /// Installs the shutdown handle for this source context.
45    ///
46    /// Called once by the runtime, before the component runs, with the shutdown signal of the component's dedicated
47    /// supervisor.
48    pub(crate) fn set_shutdown_handle(&mut self, shutdown_handle: ShutdownHandle) {
49        self.shutdown_handle = Some(shutdown_handle);
50    }
51
52    /// Installs the shutdown handle for this source context, for tests.
53    ///
54    /// The topology runtime does this itself before a source runs, using the shutdown signal of the component's
55    /// dedicated supervisor. A test that drives a source through a real shutdown has to stand in for it.
56    #[cfg(any(test, feature = "test-util"))]
57    pub fn set_shutdown_handle_for_test(&mut self, shutdown_handle: ShutdownHandle) {
58        self.set_shutdown_handle(shutdown_handle);
59    }
60
61    /// Consumes the shutdown handle of this source context.
62    ///
63    /// # Panics
64    ///
65    /// Panics if the shutdown handle has already been taken.
66    pub fn take_shutdown_handle(&mut self) -> ShutdownHandle {
67        self.shutdown_handle.take().expect("shutdown handle already taken")
68    }
69
70    /// Consumes the health handle of this source context.
71    ///
72    /// # Panics
73    ///
74    /// Panics if the health handle has already been taken.
75    pub fn take_health_handle(&mut self) -> Health {
76        self.health_handle.take().expect("health handle already taken")
77    }
78
79    /// Returns a reference to the topology context.
80    pub fn topology_context(&self) -> &TopologyContext {
81        &self.inner.topology_context
82    }
83
84    /// Returns a reference to the component context.
85    pub fn component_context(&self) -> &ComponentContext {
86        &self.inner.component_context
87    }
88
89    /// Returns a reference to the component registry.
90    pub fn component_registry(&self) -> &ComponentRegistry {
91        &self.inner.component_registry
92    }
93
94    /// Returns a reference to the events dispatcher.
95    pub fn dispatcher(&self) -> &EventsDispatcher {
96        &self.inner.dispatcher
97    }
98}
99
100impl Clone for SourceContext {
101    fn clone(&self) -> Self {
102        Self {
103            shutdown_handle: None,
104            health_handle: None,
105            inner: self.inner.clone(),
106        }
107    }
108}