1use crate::{CallContext, DbScopeDiagnosticIdentity};
4use serde::{Serialize, Serializer, ser::SerializeStruct};
5use std::sync::{
6 Arc, OnceLock,
7 atomic::{AtomicU64, Ordering},
8};
9
10static NEXT_ROOT: AtomicU64 = AtomicU64::new(1);
11
12#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
14#[serde(tag = "state", content = "value", rename_all = "snake_case")]
15pub enum ContextFact<T> {
16 Present(T),
17 NotApplicable,
18 NotEstablished,
19 Unavailable,
20}
21
22#[derive(Clone, Copy, Eq, PartialEq)]
24pub struct ContextIdentity {
25 bytes: [u8; 256],
26 len: u16,
27}
28impl ContextIdentity {
29 pub fn checked(value: &str) -> Result<Self, ContextConflict> {
30 if value.is_empty() || value.len() > 256 || value.chars().any(char::is_control) {
31 return Err(ContextConflict::InvalidIdentity);
32 }
33 let mut out = Self {
34 bytes: [0; 256],
35 len: value.len() as u16,
36 };
37 out.bytes[..value.len()].copy_from_slice(value.as_bytes());
38 Ok(out)
39 }
40 fn as_str(&self) -> &str {
41 std::str::from_utf8(&self.bytes[..usize::from(self.len)]).expect("validated UTF-8")
42 }
43}
44impl Serialize for ContextIdentity {
45 fn serialize<S: Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
46 s.serialize_str(self.as_str())
47 }
48}
49
50#[derive(Clone, Copy, Eq, PartialEq, Serialize)]
52#[serde(transparent)]
53pub struct ContextLabel(ContextIdentity);
54impl ContextLabel {
55 pub fn checked(value: &str) -> Result<Self, ContextConflict> {
56 let value = ContextIdentity::checked(value)?;
57 if value.as_str().contains("://")
58 || !value
59 .as_str()
60 .chars()
61 .all(|c| c.is_alphanumeric() || "_.:/{}*-".contains(c))
62 {
63 return Err(ContextConflict::UnsafeMetadata);
64 }
65 Ok(Self(value))
66 }
67}
68
69#[derive(Clone, Copy, Debug, Eq, PartialEq)]
70pub enum ContextConflict {
71 InvalidIdentity,
72 UnsafeMetadata,
73 IdentityGroup,
74 ForeignRoot,
75 ChildRelation,
76 CounterExhausted,
77}
78
79#[derive(Eq, PartialEq)]
82pub struct RequestIdentityGroup {
83 application: ContextLabel,
84 module: ContextLabel,
85 service: ContextLabel,
86 operation: ContextLabel,
87 trace: ContextIdentity,
88 rpc: ContextFact<ContextIdentity>,
89 span: u64,
90 request: ContextIdentity,
91 route: ContextLabel,
92 attempt: u32,
93 zone: ContextFact<ContextLabel>,
94}
95impl RequestIdentityGroup {
96 pub fn from_validated(
97 call: &CallContext,
98 request: &str,
99 route: &str,
100 attempt: u32,
101 zone: ContextFact<ContextLabel>,
102 ) -> Result<Self, ContextConflict> {
103 if attempt == 0 {
104 return Err(ContextConflict::InvalidIdentity);
105 }
106 Ok(Self {
107 application: ContextLabel::checked(call.application().as_str())?,
108 module: ContextLabel::checked(call.module().as_str())?,
109 service: ContextLabel::checked(call.service().as_str())?,
110 operation: ContextLabel::checked(call.operation().as_str())?,
111 trace: ContextIdentity::checked(call.trace_correlation_id().as_str())?,
112 rpc: match call.rpc_correlation_id() {
113 Some(id) => ContextFact::Present(ContextIdentity::checked(id.as_str())?),
114 None => ContextFact::Unavailable,
115 },
116 span: call.span_id().as_u64(),
117 request: ContextIdentity::checked(request)?,
118 route: ContextLabel::checked(route)?,
119 attempt,
120 zone,
121 })
122 }
123}
124
125struct RequestRoot {
126 local: u64,
127 application: ContextLabel,
128 initial: ContextFact<()>,
129 identity: OnceLock<RequestIdentityGroup>,
130}
131impl Drop for RequestRoot {
132 fn drop(&mut self) {
133 observe(self.local, "root_drop");
134 }
135}
136
137pub struct RequestRootPublisher {
144 root: Arc<RequestRoot>,
145}
146
147pub struct RequestRootRef {
149 root: Arc<RequestRoot>,
150}
151impl Clone for RequestRootRef {
152 fn clone(&self) -> Self {
153 observe(self.root.local, "root_share");
154 Self {
155 root: Arc::clone(&self.root),
156 }
157 }
158}
159impl Drop for RequestRootRef {
160 fn drop(&mut self) {
161 observe(self.root.local, "root_release");
162 }
163}
164
165impl RequestRootPublisher {
166 pub fn create(
169 application: ContextLabel,
170 initial: ContextFact<()>,
171 ) -> Result<Self, ContextConflict> {
172 if matches!(initial, ContextFact::Present(())) {
173 return Err(ContextConflict::InvalidIdentity);
174 }
175 let local = NEXT_ROOT
176 .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |n| n.checked_add(1))
177 .map_err(|_| ContextConflict::CounterExhausted)?;
178 let root = Arc::new(RequestRoot {
179 local,
180 application,
181 initial,
182 identity: OnceLock::new(),
183 });
184 observe(local, "root_create");
185 Ok(Self { root })
186 }
187 pub fn reference(&self) -> RequestRootRef {
188 observe(self.root.local, "root_share");
189 RequestRootRef {
190 root: Arc::clone(&self.root),
191 }
192 }
193 #[allow(clippy::result_large_err)] pub fn publish(
195 &mut self,
196 group: RequestIdentityGroup,
197 ) -> Result<(), (ContextConflict, RequestIdentityGroup)> {
198 if group.application != self.root.application {
199 return Err((ContextConflict::IdentityGroup, group));
200 }
201 if let Some(old) = self.root.identity.get() {
202 return if old == &group {
203 Ok(())
204 } else {
205 Err((ContextConflict::IdentityGroup, group))
206 };
207 }
208 self.root
209 .identity
210 .set(group)
211 .map_err(|group| (ContextConflict::IdentityGroup, group))
212 }
213}
214
215#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
217#[serde(rename_all = "snake_case")]
218pub enum RequestViewPhase {
219 SocketAccepted,
220 Reading,
221 Admitted,
222 Dispatch,
223 Handler,
224 Database,
225 Outbound,
226 Response,
227 Finalizing,
228 Finished,
229}
230
231#[derive(Clone, Copy, Eq, PartialEq, Serialize)]
233pub struct RegisteredContextOperation(&'static str);
234impl RegisteredContextOperation {
235 pub fn checked(value: &'static str) -> Result<Self, ContextConflict> {
236 ContextLabel::checked(value)?;
237 Ok(Self(value))
238 }
239}
240
241#[derive(Clone, Copy)]
242pub struct RequestLocalFacts {
243 task: ContextFact<u64>,
244 scope: ContextFact<DbScopeDiagnosticIdentity>,
245 db_operation: ContextFact<RegisteredContextOperation>,
246 phase: RequestViewPhase,
247}
248impl RequestLocalFacts {
249 pub fn new(phase: RequestViewPhase) -> Self {
250 Self {
251 task: ContextFact::NotEstablished,
252 scope: ContextFact::NotEstablished,
253 db_operation: ContextFact::NotApplicable,
254 phase,
255 }
256 }
257 pub fn with_task(mut self, task: ContextFact<u64>) -> Self {
259 self.task = task;
260 self
261 }
262 pub fn without_scope(mut self, state: ContextFact<()>) -> Self {
263 self.scope = absent(state);
264 self
265 }
266 pub fn with_db_operation(mut self, operation: RegisteredContextOperation) -> Self {
267 self.db_operation = ContextFact::Present(operation);
268 self
269 }
270}
271
272struct ChildCall {
275 #[cfg(test)]
276 observation: LocalObservation,
277 application: ContextLabel,
278 module: ContextLabel,
279 service: ContextLabel,
280 operation: ContextLabel,
281 rpc: ContextIdentity,
282 span: u64,
283 route: ContextLabel,
284 attempt: u32,
285}
286struct LocalView {
287 #[cfg(test)]
288 observation: LocalObservation,
289 root: Arc<RequestRoot>,
290 published: bool,
291 facts: RequestLocalFacts,
292 child: Option<Arc<ChildCall>>,
293}
294impl Drop for LocalView {
295 fn drop(&mut self) {
296 observe(self.root.local, "view_drop");
297 #[cfg(test)]
298 self.observation.record("destroy");
299 }
300}
301impl Drop for ChildCall {
302 fn drop(&mut self) {
303 #[cfg(test)]
304 self.observation.record("destroy");
305 }
306}
307
308pub struct RequestExecutionView {
316 inner: Arc<LocalView>,
317}
318impl Clone for RequestExecutionView {
319 fn clone(&self) -> Self {
320 #[cfg(test)]
321 self.inner.observation.record("share");
322 Self {
323 inner: Arc::clone(&self.inner),
324 }
325 }
326}
327impl Drop for RequestExecutionView {
328 fn drop(&mut self) {
329 #[cfg(test)]
330 self.inner.observation.record("release");
331 }
332}
333impl RequestRootRef {
334 pub fn view(&self, facts: RequestLocalFacts) -> RequestExecutionView {
335 RequestExecutionView::allocate(
336 Arc::clone(&self.root),
337 self.root.identity.get().is_some(),
338 facts,
339 None,
340 )
341 }
342 pub fn same_request(&self, view: &RequestExecutionView) -> bool {
343 Arc::ptr_eq(&self.root, &view.inner.root)
344 }
345}
346impl RequestExecutionView {
347 fn allocate(
348 root: Arc<RequestRoot>,
349 published: bool,
350 facts: RequestLocalFacts,
351 child: Option<Arc<ChildCall>>,
352 ) -> Self {
353 observe(root.local, "view_create");
354 Self {
355 inner: Arc::new(LocalView {
356 #[cfg(test)]
357 observation: LocalObservation::create(root.local, "view"),
358 root,
359 published,
360 facts,
361 child,
362 }),
363 }
364 }
365 fn local(&self, facts: RequestLocalFacts) -> Self {
367 Self::allocate(
368 Arc::clone(&self.inner.root),
369 self.inner.published,
370 facts,
371 self.inner.child.clone(),
372 )
373 }
374 pub fn with_phase(&self, phase: RequestViewPhase) -> Self {
375 let mut facts = self.inner.facts;
376 facts.phase = phase;
377 self.local(facts)
378 }
379 pub fn with_db_operation(&self, operation: RegisteredContextOperation) -> Self {
380 self.local(self.inner.facts.with_db_operation(operation))
381 }
382 pub fn in_task(&self, task: u64) -> Self {
385 self.local(self.inner.facts.with_task(ContextFact::Present(task)))
386 }
387 pub fn in_db_scope(
388 &self,
389 scope: &crate::DbScopeDiagnosticContext<Self>,
390 ) -> Result<Self, ContextConflict> {
391 let (original, identity) = scope.diagnostic_context();
392 if !self.same_request(original) {
393 return Err(ContextConflict::ForeignRoot);
394 }
395 let mut facts = self.inner.facts;
396 facts.scope = ContextFact::Present(identity);
397 Ok(self.local(facts))
398 }
399 pub fn observed_db_scope(
400 &self,
401 scope: &crate::DbScopeObservation<Self>,
402 ) -> Result<Self, ContextConflict> {
403 let (original, identity) = scope.diagnostic_context();
404 if !self.same_request(original) {
405 return Err(ContextConflict::ForeignRoot);
406 }
407 let mut facts = self.inner.facts;
408 facts.scope = ContextFact::Present(identity);
409 Ok(self.local(facts))
410 }
411 pub fn same_request(&self, other: &Self) -> bool {
412 Arc::ptr_eq(&self.inner.root, &other.inner.root)
413 }
414 pub fn same_view(&self, other: &Self) -> bool {
415 Arc::ptr_eq(&self.inner, &other.inner)
416 }
417 pub fn refresh(&self, root: &RequestRootRef) -> Result<Self, ContextConflict> {
419 if !root.same_request(self) {
420 return Err(ContextConflict::ForeignRoot);
421 }
422 Ok(Self::allocate(
423 Arc::clone(&self.inner.root),
424 self.inner.root.identity.get().is_some(),
425 self.inner.facts,
426 self.inner.child.clone(),
427 ))
428 }
429 pub fn child(
430 &self,
431 call: &CallContext,
432 request: &str,
433 route: &str,
434 attempt: u32,
435 ) -> Result<Self, ContextConflict> {
436 let identity = self.identity().ok_or(ContextConflict::ChildRelation)?;
437 if identity.trace.as_str() != call.trace_correlation_id().as_str()
438 || identity.request.as_str() != request
439 {
440 return Err(ContextConflict::ForeignRoot);
441 }
442 let parent_rpc = if let Some(child) = &self.inner.child {
443 &child.rpc
444 } else if let ContextFact::Present(rpc) = &identity.rpc {
445 rpc
446 } else {
447 return Err(ContextConflict::ChildRelation);
448 };
449 let rpc = call
450 .rpc_correlation_id()
451 .ok_or(ContextConflict::ChildRelation)?;
452 let suffix = rpc
453 .as_str()
454 .strip_prefix(parent_rpc.as_str())
455 .and_then(|s| s.strip_prefix('.'))
456 .ok_or(ContextConflict::ChildRelation)?;
457 let span = self
458 .inner
459 .child
460 .as_ref()
461 .map_or(identity.span, |child| child.span);
462 if suffix.is_empty()
463 || !suffix.bytes().all(|b| b.is_ascii_digit())
464 || span == call.span_id().as_u64()
465 || attempt == 0
466 {
467 return Err(ContextConflict::ChildRelation);
468 }
469 let child = Arc::new(ChildCall {
470 application: ContextLabel::checked(call.application().as_str())?,
471 module: ContextLabel::checked(call.module().as_str())?,
472 service: ContextLabel::checked(call.service().as_str())?,
473 operation: ContextLabel::checked(call.operation().as_str())?,
474 rpc: ContextIdentity::checked(rpc.as_str())?,
475 span: call.span_id().as_u64(),
476 route: ContextLabel::checked(route)?,
477 attempt,
478 #[cfg(test)]
479 observation: LocalObservation::create(self.inner.root.local, "child"),
480 });
481 Ok(Self::allocate(
482 Arc::clone(&self.inner.root),
483 self.inner.published,
484 self.inner.facts,
485 Some(child),
486 ))
487 }
488 fn identity(&self) -> Option<&RequestIdentityGroup> {
489 self.inner
490 .published
491 .then(|| self.inner.root.identity.get())
492 .flatten()
493 }
494}
495
496fn absent<T>(state: ContextFact<()>) -> ContextFact<T> {
497 match state {
498 ContextFact::NotApplicable => ContextFact::NotApplicable,
499 ContextFact::NotEstablished => ContextFact::NotEstablished,
500 _ => ContextFact::Unavailable,
501 }
502}
503impl Serialize for RequestExecutionView {
504 fn serialize<S: Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
505 let mut out = s.serialize_struct("RequestContext", 21)?;
506 let identity = self.identity();
507 let child = self.inner.child.as_deref();
508 let initial = self.inner.root.initial;
509 out.serialize_field("schema_version", &2u8)?;
510 out.serialize_field("local_request", &self.inner.root.local)?;
511 out.serialize_field("publication", &u8::from(self.inner.published))?;
512 out.serialize_field(
513 "application",
514 &ContextFact::Present(&self.inner.root.application),
515 )?;
516 out.serialize_field(
517 "call_application",
518 &ContextFact::Present(child.map_or(&self.inner.root.application, |c| &c.application)),
519 )?;
520 macro_rules! root_field {
521 ($name:literal, $field:ident) => {
522 out.serialize_field(
523 $name,
524 &identity.map_or_else(|| absent(initial), |g| ContextFact::Present(&g.$field)),
525 )?;
526 };
527 }
528 macro_rules! call_field {
529 ($name:literal, $field:ident) => {
530 out.serialize_field(
531 $name,
532 &child
533 .map(|c| &c.$field)
534 .or_else(|| identity.map(|g| &g.$field))
535 .map_or_else(|| absent(initial), ContextFact::Present),
536 )?;
537 };
538 }
539 call_field!("module", module);
540 call_field!("service", service);
541 call_field!("operation", operation);
542 root_field!("trace_id", trace);
543 root_field!("request", request);
544 let span = child.map(|c| c.span).or_else(|| identity.map(|g| g.span));
545 out.serialize_field(
546 "span_id",
547 &span.map_or_else(
548 || absent(initial),
549 |s| ContextFact::Present(SpanProjection(s)),
550 ),
551 )?;
552 call_field!("route", route);
553 call_field!("attempt", attempt);
554 let rpc = child
555 .map(|c| ContextFact::Present(c.rpc))
556 .or_else(|| identity.map(|g| g.rpc))
557 .unwrap_or_else(|| absent(initial));
558 out.serialize_field("rpc_id", &rpc)?;
559 out.serialize_field(
560 "zone",
561 &identity.map_or_else(|| absent(initial), |g| g.zone),
562 )?;
563 out.serialize_field("db_operation", &self.inner.facts.db_operation)?;
564 out.serialize_field("scope", &self.inner.facts.scope)?;
565 out.serialize_field("task", &self.inner.facts.task)?;
566 out.serialize_field("lifecycle", &ContextFact::Present(self.inner.facts.phase))?;
567 out.serialize_field(
569 "target",
570 &child.map_or(ContextFact::NotApplicable, |c| {
571 ContextFact::Present(&c.route)
572 }),
573 )?;
574 out.end()
575 }
576}
577
578pub fn request_context_layouts() -> [(std::alloc::Layout, std::alloc::Layout); 3] {
582 fn pair<T>() -> (std::alloc::Layout, std::alloc::Layout) {
583 let payload = std::alloc::Layout::new::<T>();
584 let header = std::alloc::Layout::new::<[std::sync::atomic::AtomicUsize; 2]>();
585 (
586 payload,
587 header
588 .extend(payload)
589 .expect("fixed layout")
590 .0
591 .pad_to_align(),
592 )
593 }
594 [
595 pair::<RequestRoot>(),
596 pair::<LocalView>(),
597 pair::<ChildCall>(),
598 ]
599}
600
601#[cfg(not(test))]
602fn observe(_: u64, _: &'static str) {}
603#[cfg(test)]
604fn observe(root: u64, event: &'static str) {
605 EVENTS.lock().unwrap().push((root, event));
606 record_observation(root, root, "root", event);
607}
608#[cfg(test)]
609static EVENTS: std::sync::Mutex<Vec<(u64, &'static str)>> = std::sync::Mutex::new(Vec::new());
610
611struct SpanProjection(u64);
612impl Serialize for SpanProjection {
613 fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
614 let mut bytes = [b'0'; 16];
615 for (i, byte) in bytes.iter_mut().enumerate() {
616 *byte = b"0123456789abcdef"[((self.0 >> ((15 - i) * 4)) & 15) as usize];
617 }
618 serializer.serialize_str(std::str::from_utf8(&bytes).expect("hex"))
619 }
620}
621
622#[cfg(test)]
625struct LocalObservation {
626 id: u64,
627 root: u64,
628 kind: &'static str,
629}
630#[cfg(test)]
631static LOCAL_EVENTS: std::sync::Mutex<Vec<(u64, u64, &'static str, &'static str)>> =
632 std::sync::Mutex::new(Vec::new());
633#[cfg(test)]
634impl LocalObservation {
635 fn create(root: u64, kind: &'static str) -> Self {
636 static NEXT: AtomicU64 = AtomicU64::new(1);
637 let observation = Self {
638 id: NEXT.fetch_add(1, Ordering::Relaxed),
639 root,
640 kind,
641 };
642 observation.record("create");
643 observation
644 }
645 fn record(&self, event: &'static str) {
646 LOCAL_EVENTS
647 .lock()
648 .unwrap()
649 .push((self.id, self.root, self.kind, event));
650 record_observation(self.root, self.id, self.kind, event);
651 }
652}
653
654#[cfg(test)]
655type ObservationEvent = (u64, u64, u64, &'static str, &'static str);
656#[cfg(test)]
657static ORDERED_EVENTS: std::sync::Mutex<Vec<ObservationEvent>> = std::sync::Mutex::new(Vec::new());
658#[cfg(test)]
659fn record_observation(root: u64, object: u64, kind: &'static str, event: &'static str) {
660 static SEQUENCE: AtomicU64 = AtomicU64::new(1);
661 let sequence = SEQUENCE.fetch_add(1, Ordering::Relaxed);
662 ORDERED_EVENTS
663 .lock()
664 .unwrap()
665 .push((sequence, root, object, kind, event));
666}
667
668#[cfg(test)]
669mod tests;