dcontext_dactor/inbound.rs
1use std::any::Any;
2
3use dactor::{
4 Disposition, HandlerWrapper, Headers, InboundContext, InboundInterceptor, Outcome,
5 RuntimeHeaders,
6};
7
8use crate::header::{ContextHeader, ContextSnapshotHeader};
9use crate::propagation::bytes_to_snapshot;
10use crate::ErrorPolicy;
11
12/// Inbound interceptor that propagates dcontext automatically through actor
13/// message handlers.
14///
15/// This interceptor performs two complementary actions:
16///
17/// 1. **`on_receive`** — Normalizes context headers: if only wire bytes
18/// ([`ContextHeader`]) are present (remote hop), deserializes them into a
19/// [`ContextSnapshotHeader`]. Local snapshots are left as-is.
20///
21/// 2. **`wrap_handler`** — Wraps the handler future with
22/// [`dcontext::with_context`], restoring the propagated context into the
23/// async task-local scope. This makes `dcontext::get_context` /
24/// `dcontext::set_context` work automatically inside the handler — **no
25/// manual `with_propagated_context()` call needed**.
26///
27/// ## Error Handling
28///
29/// Controlled by [`ErrorPolicy`]:
30/// - [`LogAndContinue`](ErrorPolicy::LogAndContinue) (default) — log and
31/// deliver the message without context.
32/// - [`Reject`](ErrorPolicy::Reject) — reject the message.
33///
34/// # Usage
35///
36/// ```ignore
37/// use dcontext_dactor::{ContextOutboundInterceptor, ContextInboundInterceptor};
38///
39/// // Register interceptors — context propagation is fully automatic
40/// runtime.add_outbound_interceptor(Box::new(ContextOutboundInterceptor::default()));
41/// runtime.add_inbound_interceptor(Box::new(ContextInboundInterceptor::default()));
42///
43/// // In the handler, dcontext is available without any boilerplate:
44/// #[async_trait]
45/// impl Handler<MyMessage> for MyActor {
46/// async fn handle(&mut self, msg: MyMessage, ctx: &mut ActorContext) -> () {
47/// let rid: RequestId = dcontext::get_context("request_id");
48/// // ... context is automatically restored by the interceptor
49/// }
50/// }
51/// ```
52pub struct ContextInboundInterceptor {
53 error_policy: ErrorPolicy,
54}
55
56impl ContextInboundInterceptor {
57 /// Create with a specific error policy.
58 pub fn new(error_policy: ErrorPolicy) -> Self {
59 Self { error_policy }
60 }
61}
62
63impl Default for ContextInboundInterceptor {
64 fn default() -> Self {
65 Self {
66 error_policy: ErrorPolicy::default(),
67 }
68 }
69}
70
71impl InboundInterceptor for ContextInboundInterceptor {
72 fn name(&self) -> &'static str {
73 "dcontext-inbound"
74 }
75
76 fn on_receive(
77 &self,
78 _ctx: &InboundContext<'_>,
79 _runtime_headers: &RuntimeHeaders,
80 headers: &mut Headers,
81 _message: &dyn Any,
82 ) -> Disposition {
83 // If a local snapshot header already exists, nothing to do.
84 if headers.get::<ContextSnapshotHeader>().is_some() {
85 return Disposition::Continue;
86 }
87
88 // Try to convert wire bytes into a snapshot for uniform handler access.
89 if let Some(wire_header) = headers.get::<ContextHeader>() {
90 let bytes = &wire_header.bytes;
91 match bytes_to_snapshot(bytes) {
92 Some(snap) => {
93 headers.insert(ContextSnapshotHeader { snapshot: snap });
94 }
95 None => {
96 tracing::warn!(
97 target: "dcontext_dactor",
98 "failed to deserialize dcontext from inbound wire bytes"
99 );
100 if self.error_policy == ErrorPolicy::Reject {
101 return Disposition::Reject(
102 "dcontext deserialization failed: corrupt wire bytes".into(),
103 );
104 }
105 }
106 }
107 }
108
109 Disposition::Continue
110 }
111
112 fn wrap_handler<'a>(
113 &'a self,
114 ctx: &InboundContext<'_>,
115 headers: &Headers,
116 ) -> Option<HandlerWrapper<'a>> {
117 // Extract the snapshot prepared by on_receive (or outbound for local hops).
118 let snapshot = headers.get::<ContextSnapshotHeader>()?.snapshot.clone();
119 // Create a named scope for the remote call boundary, e.g. "remote:MyActor".
120 let scope_name = format!("remote:{}", ctx.actor_name);
121 Some(Box::new(move |next| {
122 Box::pin(async move {
123 let result = dcontext::with_context(snapshot, async move {
124 dcontext::named_scope_async(scope_name, next).await
125 }).await;
126 result
127 })
128 }))
129 }
130
131 fn on_complete(
132 &self,
133 _ctx: &InboundContext<'_>,
134 _runtime_headers: &RuntimeHeaders,
135 _headers: &Headers,
136 _outcome: &Outcome<'_>,
137 ) {
138 // No cleanup needed — context scope ends when the wrapped future completes.
139 }
140}