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#[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 pub const ROOT: Self = Self(0);
33
34 pub fn new() -> Self {
36 let id = GLOBAL_PROCESS_ID_COUNTER.fetch_add(1, Relaxed);
37 Self(id)
38 }
39
40 pub fn current() -> Self {
44 CURRENT_PROCESS_ID.try_with(|id| *id).unwrap_or(Self::ROOT)
45 }
46
47 pub fn as_usize(&self) -> usize {
49 self.0
50 }
51}
52
53#[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#[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_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 pub fn id(&self) -> &Id {
160 &self.id
161 }
162
163 pub(crate) fn name(&self) -> &str {
166 &self.name
167 }
168
169 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#[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
239pub trait ProcessExt {
241 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 if name.is_empty() {
262 return false;
263 }
264
265 if !name.starts_with(|c: char| c.is_alphanumeric()) || !name.ends_with(|c: char| c.is_alphanumeric()) {
267 return false;
268 }
269
270 for c in name.chars() {
274 if !c.is_alphanumeric() && c != '_' {
275 return false;
276 }
277 }
278
279 true
280}
281
282pub fn get_sanitized_name(name: &str) -> MetaString {
300 if is_process_name_segment_valid(name) {
301 name.into()
302 } else {
303 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 let trimmed = sanitized.trim_matches(|c: char| !c.is_alphanumeric());
324 trimmed.into()
325 }
326}
327
328fn 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}