dcontext_dactor/
inbound.rs

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

dcontext_dactor/
inbound.rs
