Skip to main content

saluki_core/runtime/
process.rs

1use std::{
2    future::Future,
3    ops::Deref,
4    pin::Pin,
5    sync::atomic::{AtomicUsize, Ordering::Relaxed},
6    task::{Context, Poll},
7};
8
9use pin_project::{pin_project, pinned_drop};
10use saluki_common::resource_tracking::{ResourceGroupRegistry, ResourceGroupToken, Track as _, Tracked};
11use stringtheory::MetaString;
12use tracing::{debug_span, instrument::Instrumented, Instrument as _};
13
14use super::state::{DataspaceRegistry, CURRENT_DATASPACE};
15
16static GLOBAL_PROCESS_ID_COUNTER: AtomicUsize = AtomicUsize::new(1);
17
18/// Process identifier.
19///
20/// A simple, numeric identifier that uniquely identifies a process.
21///
22/// Guaranteed to be unique.
23#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
24pub struct Id(usize);
25
26tokio::task_local! {
27    pub(crate) static CURRENT_PROCESS_ID: Id;
28}
29
30impl Id {
31    /// The root process identifier, representing the global/unnamed process context.
32    pub const ROOT: Self = Self(0);
33
34    /// Creates a new process identifier.
35    pub fn new() -> Self {
36        let id = GLOBAL_PROCESS_ID_COUNTER.fetch_add(1, Relaxed);
37        Self(id)
38    }
39
40    /// Returns the process identifier for the currently executing process.
41    ///
42    /// If called outside of a process context, returns `Id::ROOT`.
43    pub fn current() -> Self {
44        CURRENT_PROCESS_ID.try_with(|id| *id).unwrap_or(Self::ROOT)
45    }
46
47    /// Returns the raw numeric value of this identifier.
48    pub fn as_usize(&self) -> usize {
49        self.0
50    }
51}
52
53/// Process name.
54///
55/// A human-readable name for a process that only contains alphanumeric characters, underscores, and periods.
56///
57/// Process names are scoped, such that the resulting process name is nested. For example, if a supervisor has a name of
58/// `topology_sup`, and a child process is added to that supervisor with a name of `worker`, the resulting process name
59/// for the child process will be `topology_sup.worker`. Process names can be arbitrarily nested in this way.
60///
61/// Process names will be sanitized if they contain invalid characters, such as hyphens or spaces. Invalid characters
62/// will be replaced with underscores.
63///
64/// A period in a provided name is treated as a segment separator: the name is split on periods, each segment is
65/// sanitized independently, and the segments are rejoined with periods. This lets a caller supply an already-scoped
66/// name (such as `topology.primary`) and have it preserved as distinct segments rather than collapsed.
67///
68/// Not guaranteed to be unique.
69#[derive(Clone, Debug, PartialEq, Eq, Hash)]
70pub struct Name(MetaString);
71
72impl Name {
73    pub(crate) fn root<N: AsRef<str>>(name: N) -> Option<Self> {
74        sanitize_scoped_name(name.as_ref()).map(Self)
75    }
76
77    pub(crate) fn scoped<N: AsRef<str>>(parent: &Name, name: N) -> Option<Self> {
78        let child = sanitize_scoped_name(name.as_ref())?;
79        Some(Self(format!("{}.{}", parent.0, child).into()))
80    }
81}
82
83impl Deref for Name {
84    type Target = str;
85
86    fn deref(&self) -> &Self::Target {
87        &self.0
88    }
89}
90
91/// A runtime process.
92#[derive(Clone)]
93pub struct Process {
94    id: Id,
95    name: Name,
96    resource_group_token: ResourceGroupToken,
97    dataspace: DataspaceRegistry,
98}
99
100impl Process {
101    pub(crate) fn supervisor<N: AsRef<str>>(name: N, parent: Option<&Process>) -> Option<Self> {
102        let name = parent
103            .and_then(|p| Name::scoped(&p.name, &name))
104            .or_else(|| Name::root(name))?;
105        let resource_group_token = ResourceGroupRegistry::global().register_resource_group(&*name);
106        let dataspace = parent.map(|p| p.dataspace.clone()).unwrap_or_default();
107        Some(Self::from_parts(Id::new(), name, resource_group_token, dataspace))
108    }
109
110    pub(crate) fn supervisor_with_dataspace<N: AsRef<str>>(
111        name: N, parent: Option<&Process>, dataspace: Option<DataspaceRegistry>,
112    ) -> Option<Self> {
113        let name = parent
114            .and_then(|p| Name::scoped(&p.name, &name))
115            .or_else(|| Name::root(name))?;
116        let resource_group_token = ResourceGroupRegistry::global().register_resource_group(&*name);
117        let dataspace = dataspace
118            .or_else(|| parent.map(|p| p.dataspace.clone()))
119            .unwrap_or_default();
120        Some(Self::from_parts(Id::new(), name, resource_group_token, dataspace))
121    }
122
123    /// Creates a worker process that has its own resource group, with the same name as the process.
124    ///
125    /// The process does not inherit the resource group of its parent. Registration is idempotent by name. Thus, if the
126    /// name of a worker matches a group registered elsewhere, the worker shares that group. An example is the group of
127    /// a topology component, which its registry registers under the same name.
128    pub(crate) fn worker_with_own_resource_group<N: AsRef<str>>(name: N, parent: &Process) -> Option<Self> {
129        let name = Name::scoped(&parent.name, name)?;
130        let resource_group_token = ResourceGroupRegistry::global().register_resource_group(&*name);
131        Some(Self::from_parts(
132            Id::new(),
133            name,
134            resource_group_token,
135            parent.dataspace.clone(),
136        ))
137    }
138
139    pub(crate) fn worker<N: AsRef<str>>(name: N, parent: &Process) -> Option<Self> {
140        let name = Name::scoped(&parent.name, name)?;
141        Some(Self::from_parts(
142            Id::new(),
143            name,
144            parent.resource_group_token,
145            parent.dataspace.clone(),
146        ))
147    }
148
149    fn from_parts(id: Id, name: Name, resource_group_token: ResourceGroupToken, dataspace: DataspaceRegistry) -> Self {
150        Self {
151            id,
152            name,
153            resource_group_token,
154            dataspace,
155        }
156    }
157
158    /// Returns the process identifier.
159    pub fn id(&self) -> &Id {
160        &self.id
161    }
162
163    /// Returns the fully qualified process name (scoped under its parent, for example
164    /// `topology.primary.sources.dsd_in.listener_udp`).
165    pub(crate) fn name(&self) -> &str {
166        &self.name
167    }
168
169    /// Returns the dataspace registry associated with this process.
170    pub(crate) fn dataspace(&self) -> &DataspaceRegistry {
171        &self.dataspace
172    }
173
174    pub fn into_process_future<F>(self, inner: F) -> ProcessFuture<F>
175    where
176        F: Future,
177    {
178        ProcessFuture::new(self, inner)
179    }
180}
181
182/// A process future.
183///
184/// Wraps a [`Future`] with process-specific instrumentation and globals. This ensures that processes have properly scoped tracing and
185/// allocation tracking behavior, as well as access to supervisor/runtime-specific globals, such as the dataspace registry.
186#[pin_project(PinnedDrop)]
187pub struct ProcessFuture<F> {
188    process_id: Id,
189    dataspace: DataspaceRegistry,
190    #[pin]
191    inner: Instrumented<Tracked<F>>,
192}
193
194impl<F> ProcessFuture<F>
195where
196    F: Future,
197{
198    pub(crate) fn new(process: Process, inner: F) -> Self {
199        let span = debug_span!(
200            "process",
201            process_id = process.id().as_usize(),
202            process_name = &*process.name,
203        );
204
205        let process_id = process.id;
206        let dataspace = process.dataspace.clone();
207        let inner = inner.track_resources(process.resource_group_token).instrument(span);
208
209        Self {
210            process_id,
211            dataspace,
212            inner,
213        }
214    }
215}
216
217impl<F> Future for ProcessFuture<F>
218where
219    F: Future,
220{
221    type Output = F::Output;
222
223    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
224        let this = self.project();
225        CURRENT_PROCESS_ID.sync_scope(*this.process_id, || {
226            CURRENT_DATASPACE.sync_scope(this.dataspace.clone(), || this.inner.poll(cx))
227        })
228    }
229}
230
231#[pinned_drop]
232impl<F> PinnedDrop for ProcessFuture<F> {
233    fn drop(self: Pin<&mut Self>) {
234        let this = self.project();
235        this.dataspace.retract_all_for_process(*this.process_id);
236    }
237}
238
239/// Helper trait for running process futures.
240pub trait ProcessExt {
241    /// Converts the future into a process future.
242    fn into_process_future(self, process: Process) -> ProcessFuture<Self>
243    where
244        Self: Future + Sized;
245}
246
247impl<F> ProcessExt for F
248where
249    F: Future,
250{
251    fn into_process_future(self, process: Process) -> ProcessFuture<Self>
252    where
253        Self: Future + Sized,
254    {
255        process.into_process_future(self)
256    }
257}
258
259fn is_process_name_segment_valid(name: &str) -> bool {
260    // Process name cannot be empty.
261    if name.is_empty() {
262        return false;
263    }
264
265    // Process names cannot start or end with anything other than alphanumeric characters.
266    if !name.starts_with(|c: char| c.is_alphanumeric()) || !name.ends_with(|c: char| c.is_alphanumeric()) {
267        return false;
268    }
269
270    // Process name segments can only include alphanumeric characters and underscores.
271    //
272    // Periods are allowed in process names overall, but they're only used as separators between segments.
273    for c in name.chars() {
274        if !c.is_alphanumeric() && c != '_' {
275            return false;
276        }
277    }
278
279    true
280}
281
282/// Sanitizes `name` into a single process-name segment.
283///
284/// Invalid characters are replaced with underscores, runs of underscores are collapsed into one, and leading and
285/// trailing non-alphanumeric characters are trimmed. A name that is already a valid segment is returned unchanged.
286///
287/// Periods are *not* treated as separators here, unlike in a scoped process name: they are invalid within a segment,
288/// so they become underscores like any other invalid character. That is what makes this the right way to fold a
289/// dotted value -- a hostname, a URL authority -- into one segment of a task name. Interpolating such a value
290/// directly would instead split it across levels of the process tree.
291///
292/// # Examples
293///
294/// ```
295/// # use saluki_core::runtime::get_sanitized_name;
296/// assert_eq!(&*get_sanitized_name("app.datadoghq.com"), "app_datadoghq_com");
297/// assert_eq!(&*get_sanitized_name("already_valid"), "already_valid");
298/// ```
299pub fn get_sanitized_name(name: &str) -> MetaString {
300    if is_process_name_segment_valid(name) {
301        name.into()
302    } else {
303        // Replace invalid characters with underscores, and collapses multiple underscores into a single one.
304        let raw_sanitized = name
305            .chars()
306            .map(|c| if c.is_alphanumeric() || c == '_' { c } else { '_' });
307        let mut sanitized = String::with_capacity(name.len());
308
309        let mut last_was_underscore = true;
310        for c in raw_sanitized {
311            if c == '_' {
312                if !last_was_underscore {
313                    sanitized.push(c);
314                    last_was_underscore = true;
315                }
316            } else {
317                sanitized.push(c);
318                last_was_underscore = false;
319            }
320        }
321
322        // Remove all non-alphanumeric characters from beginning and end.
323        let trimmed = sanitized.trim_matches(|c: char| !c.is_alphanumeric());
324        trimmed.into()
325    }
326}
327
328/// Sanitizes a (possibly dotted) name into a dotted string of process-safe segments.
329///
330/// Periods are treated as segment separators: the input is split on `.`, each segment is sanitized via
331/// [`get_sanitized_name`], and the results are rejoined with `.`. Segments that are empty (or sanitize to empty) are
332/// dropped. Returns `None` if nothing remains.
333fn sanitize_scoped_name(name: &str) -> Option<MetaString> {
334    let mut rendered = String::new();
335    for segment in name.split('.') {
336        let sanitized = get_sanitized_name(segment);
337        if sanitized.is_empty() {
338            continue;
339        }
340
341        if !rendered.is_empty() {
342            rendered.push('.');
343        }
344        rendered.push_str(&sanitized);
345    }
346
347    if rendered.is_empty() {
348        None
349    } else {
350        Some(rendered.into())
351    }
352}
353
354#[cfg(test)]
355mod tests {
356    use super::*;
357
358    #[test]
359    fn test_process_name_root() {
360        let cases = [
361            ("topology_sup", Some("topology_sup")),
362            ("worker", Some("worker")),
363            ("worker.", Some("worker")),
364            ("_worker_", Some("worker")),
365            ("worker-123", Some("worker_123")),
366            ("--worker_123", Some("worker_123")),
367            ("worker 123", Some("worker_123")),
368            ("worker===123", Some("worker_123")),
369            ("topology.worker", Some("topology.worker")),
370            ("", None),
371        ];
372
373        for (input, expected) in cases {
374            let name = Name::root(input);
375            assert_eq!(name.as_deref(), expected);
376        }
377    }
378
379    #[test]
380    fn test_process_name_scoped() {
381        let parent = Name::root("topology_sup").unwrap();
382        let cases = [
383            ("worker", Some("topology_sup.worker")),
384            ("worker.", Some("topology_sup.worker")),
385            ("_worker_", Some("topology_sup.worker")),
386            ("worker-123", Some("topology_sup.worker_123")),
387            ("--worker_123", Some("topology_sup.worker_123")),
388            ("worker 123", Some("topology_sup.worker_123")),
389            ("worker===123", Some("topology_sup.worker_123")),
390            ("nested.worker", Some("topology_sup.nested.worker")),
391            ("", None),
392        ];
393
394        for (input, expected) in cases {
395            let name = Name::scoped(&parent, input);
396            assert_eq!(name.as_deref(), expected);
397        }
398    }
399}