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