Skip to main content

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}