es_entity/context/
with_event_context.rs

1use pin_project::pin_project;
2
3use std::{
4    future::Future,
5    pin::Pin,
6    task::{Context, Poll},
7};
8
9use super::{ContextData, EventContext};
10
11/// Extension trait for propagating event context across async boundaries.
12///
13/// This trait is automatically implemented for all `Future` types and provides
14/// the `with_event_context` method to carry context data across async operations
15/// like `tokio::spawn`, `tokio::task::yield_now()`, and other async boundaries
16/// where thread-local storage is not preserved.
17///
18/// # Examples
19///
20/// ```rust
21/// use es_entity::context::{EventContext, WithEventContext};
22///
23/// async fn example() {
24///     let mut ctx = EventContext::current();
25///     ctx.insert("request_id", &"abc123").unwrap();
26///
27///     let data = ctx.data();
28///     tokio::spawn(async {
29///         // Context is available here
30///         let current = EventContext::current();
31///         // current now has the request_id from the parent
32///     }.with_event_context(data)).await.unwrap();
33/// }
34/// ```
35pub trait WithEventContext: Future {
36    /// Wraps this future with event context data.
37    ///
38    /// This method ensures that when the future is polled, the provided context
39    /// data will be available as the current event context. This is essential
40    /// for maintaining context across async boundaries where the original
41    /// thread-local context is not available.
42    ///
43    /// # Arguments
44    ///
45    /// * `context_data` - The context data to make available during future execution
46    ///
47    /// # Returns
48    ///
49    /// Returns an [`EventContextFuture`] that will poll the wrapped future with
50    /// the provided context active.
51    fn with_event_context(self, context_data: ContextData) -> EventContextFuture<Self>
52    where
53        Self: Sized,
54    {
55        EventContextFuture {
56            future: self,
57            context_data,
58        }
59    }
60}
61
62impl<F: Future> WithEventContext for F {}
63
64/// A future wrapper that provides event context during polling.
65///
66/// This struct is created by the `with_event_context` method and should not
67/// be constructed directly. It ensures that the wrapped future has access
68/// to the specified event context data whenever it is polled.
69///
70/// The future maintains context isolation - the context is only active
71/// during the polling of the wrapped future and does not leak to other
72/// concurrent operations.
73#[pin_project]
74pub struct EventContextFuture<F> {
75    #[pin]
76    future: F,
77    context_data: ContextData,
78}
79
80impl<F: Future> Future for EventContextFuture<F> {
81    type Output = F::Output;
82
83    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
84        let this = self.project();
85        let ctx = EventContext::seed(this.context_data.clone());
86        let res = this.future.poll(cx);
87        *this.context_data = ctx.data();
88        res
89    }
90}
es_entity/context/with_event_context.rs

es_entity/context/
with_event_context.rs
