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    pub(crate) fn worker<N: AsRef<str>>(name: N, parent: &Process) -> Option<Self> {
124        let name = Name::scoped(&parent.name, name)?;
125        Some(Self::from_parts(
126            Id::new(),
127            name,
128            parent.resource_group_token,
129            parent.dataspace.clone(),
130        ))
131    }
132
133    fn from_parts(id: Id, name: Name, resource_group_token: ResourceGroupToken, dataspace: DataspaceRegistry) -> Self {
134        Self {
135            id,
136            name,
137            resource_group_token,
138            dataspace,
139        }
140    }
141
142    /// Returns the process identifier.
143    pub fn id(&self) -> &Id {
144        &self.id
145    }
146
147    /// Returns the fully qualified process name (scoped under its parent, for example
148    /// `topology.primary.sources.dsd_in.source`).
149    pub(crate) fn name(&self) -> &str {
150        &self.name
151    }
152
153    /// Returns the dataspace registry associated with this process.
154    pub(crate) fn dataspace(&self) -> &DataspaceRegistry {
155        &self.dataspace
156    }
157
158    pub fn into_process_future<F>(self, inner: F) -> ProcessFuture<F>
159    where
160        F: Future,
161    {
162        ProcessFuture::new(self, inner)
163    }
164}
165
166/// A process future.
167///
168/// Wraps a [`Future`] with process-specific instrumentation and globals. This ensures that processes have properly scoped tracing and
169/// allocation tracking behavior, as well as access to supervisor/runtime-specific globals, such as the dataspace registry.
170#[pin_project(PinnedDrop)]
171pub struct ProcessFuture<F> {
172    process_id: Id,
173    dataspace: DataspaceRegistry,
174    #[pin]
175    inner: Instrumented<Tracked<F>>,
176}
177
178impl<F> ProcessFuture<F>
179where
180    F: Future,
181{
182    pub(crate) fn new(process: Process, inner: F) -> Self {
183        let span = debug_span!(
184            "process",
185            process_id = process.id().as_usize(),
186            process_name = &*process.name,
187        );
188
189        let process_id = process.id;
190        let dataspace = process.dataspace.clone();
191        let inner = inner.track_resources(process.resource_group_token).instrument(span);
192
193        Self {
194            process_id,
195            dataspace,
196            inner,
197        }
198    }
199}
200
201impl<F> Future for ProcessFuture<F>
202where
203    F: Future,
204{
205    type Output = F::Output;
206
207    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
208        let this = self.project();
209        CURRENT_PROCESS_ID.sync_scope(*this.process_id, || {
210            CURRENT_DATASPACE.sync_scope(this.dataspace.clone(), || this.inner.poll(cx))
211        })
212    }
213}
214
215#[pinned_drop]
216impl<F> PinnedDrop for ProcessFuture<F> {
217    fn drop(self: Pin<&mut Self>) {
218        let this = self.project();
219        this.dataspace.retract_all_for_process(*this.process_id);
220    }
221}
222
223/// Helper trait for running process futures.
224pub trait ProcessExt {
225    /// Converts the future into a process future.
226    fn into_process_future(self, process: Process) -> ProcessFuture<Self>
227    where
228        Self: Future + Sized;
229}
230
231impl<F> ProcessExt for F
232where
233    F: Future,
234{
235    fn into_process_future(self, process: Process) -> ProcessFuture<Self>
236    where
237        Self: Future + Sized,
238    {
239        process.into_process_future(self)
240    }
241}
242
243fn is_process_name_segment_valid(name: &str) -> bool {
244    // Process name cannot be empty.
245    if name.is_empty() {
246        return false;
247    }
248
249    // Process names cannot start or end with anything other than alphanumeric characters.
250    if !name.starts_with(|c: char| c.is_alphanumeric()) || !name.ends_with(|c: char| c.is_alphanumeric()) {
251        return false;
252    }
253
254    // Process name segments can only include alphanumeric characters and underscores.
255    //
256    // Periods are allowed in process names overall, but they're only used as separators between segments.
257    for c in name.chars() {
258        if !c.is_alphanumeric() && c != '_' {
259            return false;
260        }
261    }
262
263    true
264}
265
266pub(crate) fn get_sanitized_name(name: &str) -> MetaString {
267    if is_process_name_segment_valid(name) {
268        name.into()
269    } else {
270        // Replace invalid characters with underscores, and collapses multiple underscores into a single one.
271        let raw_sanitized = name
272            .chars()
273            .map(|c| if c.is_alphanumeric() || c == '_' { c } else { '_' });
274        let mut sanitized = String::with_capacity(name.len());
275
276        let mut last_was_underscore = true;
277        for c in raw_sanitized {
278            if c == '_' {
279                if !last_was_underscore {
280                    sanitized.push(c);
281                    last_was_underscore = true;
282                }
283            } else {
284                sanitized.push(c);
285                last_was_underscore = false;
286            }
287        }
288
289        // Remove all non-alphanumeric characters from beginning and end.
290        let trimmed = sanitized.trim_matches(|c: char| !c.is_alphanumeric());
291        trimmed.into()
292    }
293}
294
295/// Sanitizes a (possibly dotted) name into a dotted string of process-safe segments.
296///
297/// Periods are treated as segment separators: the input is split on `.`, each segment is sanitized via
298/// [`get_sanitized_name`], and the results are rejoined with `.`. Segments that are empty (or sanitize to empty) are
299/// dropped. Returns `None` if nothing remains.
300fn sanitize_scoped_name(name: &str) -> Option<MetaString> {
301    let mut rendered = String::new();
302    for segment in name.split('.') {
303        let sanitized = get_sanitized_name(segment);
304        if sanitized.is_empty() {
305            continue;
306        }
307
308        if !rendered.is_empty() {
309            rendered.push('.');
310        }
311        rendered.push_str(&sanitized);
312    }
313
314    if rendered.is_empty() {
315        None
316    } else {
317        Some(rendered.into())
318    }
319}
320
321#[cfg(test)]
322mod tests {
323    use super::*;
324
325    #[test]
326    fn test_process_name_root() {
327        let cases = [
328            ("topology_sup", Some("topology_sup")),
329            ("worker", Some("worker")),
330            ("worker.", Some("worker")),
331            ("_worker_", Some("worker")),
332            ("worker-123", Some("worker_123")),
333            ("--worker_123", Some("worker_123")),
334            ("worker 123", Some("worker_123")),
335            ("worker===123", Some("worker_123")),
336            ("topology.worker", Some("topology.worker")),
337            ("", None),
338        ];
339
340        for (input, expected) in cases {
341            let name = Name::root(input);
342            assert_eq!(name.as_deref(), expected);
343        }
344    }
345
346    #[test]
347    fn test_process_name_scoped() {
348        let parent = Name::root("topology_sup").unwrap();
349        let cases = [
350            ("worker", Some("topology_sup.worker")),
351            ("worker.", Some("topology_sup.worker")),
352            ("_worker_", Some("topology_sup.worker")),
353            ("worker-123", Some("topology_sup.worker_123")),
354            ("--worker_123", Some("topology_sup.worker_123")),
355            ("worker 123", Some("topology_sup.worker_123")),
356            ("worker===123", Some("topology_sup.worker_123")),
357            ("nested.worker", Some("topology_sup.nested.worker")),
358            ("", None),
359        ];
360
361        for (input, expected) in cases {
362            let name = Name::scoped(&parent, input);
363            assert_eq!(name.as_deref(), expected);
364        }
365    }
366}