1use std::sync::Arc;
12
13use crate::task::TaskMetadata;
14
15#[derive(Clone, Debug)]
21pub struct TaskExecutionContext {
22 pub workflow_id: Arc<str>,
29 pub instance_id: Arc<str>,
31 pub task_id: Arc<str>,
33 pub metadata: TaskMetadata,
35 pub workflow_metadata_json: Option<Arc<str>>,
37}
38
39pub struct WorkflowContext<C, M> {
46 pub workflow_id: crate::WorkflowId,
48 pub workflow_name: Arc<str>,
50 pub codec: Arc<C>,
52 pub metadata: Arc<M>,
54 pub metadata_json: Option<Arc<str>>,
56}
57
58impl<C, M> Clone for WorkflowContext<C, M> {
59 fn clone(&self) -> Self {
60 Self {
61 workflow_id: self.workflow_id,
62 workflow_name: Arc::clone(&self.workflow_name),
63 codec: Arc::clone(&self.codec),
64 metadata: Arc::clone(&self.metadata),
65 metadata_json: self.metadata_json.clone(),
66 }
67 }
68}
69
70impl<C, M> WorkflowContext<C, M> {
71 pub fn new(workflow_name: impl Into<Arc<str>>, codec: Arc<C>, metadata: Arc<M>) -> Self {
73 let workflow_name: Arc<str> = workflow_name.into();
74 let workflow_id = crate::WorkflowId::from(workflow_name.as_ref());
75 Self {
76 workflow_id,
77 workflow_name,
78 codec,
79 metadata,
80 metadata_json: None,
81 }
82 }
83
84 #[must_use]
86 pub fn workflow_id(&self) -> &str {
87 &self.workflow_name
88 }
89
90 #[must_use]
92 pub fn workflow_id_hash(&self) -> crate::WorkflowId {
93 self.workflow_id
94 }
95
96 #[must_use]
98 pub fn codec(&self) -> Arc<C> {
99 self.codec.clone()
100 }
101
102 #[must_use]
104 pub fn metadata(&self) -> Arc<M> {
105 self.metadata.clone()
106 }
107}
108
109use std::cell::RefCell;
110
111std::thread_local! {
112 static THREAD_LOCAL_TASK_CTX: RefCell<Option<TaskExecutionContext>> = const { RefCell::new(None) };
117}
118
119pub fn with_thread_local_task_context<R>(ctx: TaskExecutionContext, f: impl FnOnce() -> R) -> R {
122 THREAD_LOCAL_TASK_CTX.with(|cell| {
123 let prev = cell.borrow_mut().replace(ctx);
124 let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(f));
125 *cell.borrow_mut() = prev;
126 match result {
127 Ok(r) => r,
128 Err(e) => std::panic::resume_unwind(e),
129 }
130 })
131}
132
133#[must_use]
135pub fn get_thread_local_task_context() -> Option<TaskExecutionContext> {
136 THREAD_LOCAL_TASK_CTX.with(|cell| cell.borrow().clone())
137}
138
139#[cfg(feature = "tokio")]
142mod task_local_ctx {
143 use super::TaskExecutionContext;
144
145 tokio::task_local! {
146 static TASK_EXEC_CTX: Option<TaskExecutionContext>;
148 }
149
150 pub async fn with_task_context<F: std::future::Future>(
152 ctx: TaskExecutionContext,
153 fut: F,
154 ) -> F::Output {
155 TASK_EXEC_CTX.scope(Some(ctx), fut).await
156 }
157
158 #[must_use]
162 pub fn get_task_context() -> Option<TaskExecutionContext> {
163 TASK_EXEC_CTX
164 .try_with(std::clone::Clone::clone)
165 .ok()
166 .flatten()
167 .or_else(super::get_thread_local_task_context)
168 }
169}
170
171#[cfg(feature = "tokio")]
172pub use task_local_ctx::{get_task_context, with_task_context};
173
174#[cfg(not(feature = "tokio"))]
178#[must_use]
179pub fn get_task_context() -> Option<TaskExecutionContext> {
180 get_thread_local_task_context()
181}
182
183#[macro_export]
195macro_rules! task_context {
196 () => {
197 $crate::context::get_task_context()
198 };
199}
200
201#[cfg(all(test, feature = "tokio"))]
202#[allow(clippy::unwrap_used, clippy::panic)]
203mod tests {
204 use super::*;
205 use crate::task::TaskMetadata;
206
207 fn make_task_ctx() -> TaskExecutionContext {
208 TaskExecutionContext {
209 workflow_id: Arc::from("wf-1"),
210 instance_id: Arc::from("inst-1"),
211 task_id: Arc::from("task-a"),
212 metadata: TaskMetadata::default(),
213 workflow_metadata_json: None,
214 }
215 }
216
217 #[test]
218 fn thread_local_roundtrip() {
219 assert!(get_thread_local_task_context().is_none());
220
221 let ctx = make_task_ctx();
222 let result = with_thread_local_task_context(ctx.clone(), || {
223 let inner = get_thread_local_task_context().unwrap();
224 assert_eq!(&*inner.workflow_id, "wf-1");
225 assert_eq!(&*inner.instance_id, "inst-1");
226 assert_eq!(&*inner.task_id, "task-a");
227 42
228 });
229 assert_eq!(result, 42);
230
231 assert!(get_thread_local_task_context().is_none());
233 }
234
235 #[test]
236 fn thread_local_restores_on_panic() {
237 let ctx = make_task_ctx();
238 let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
239 with_thread_local_task_context(ctx, || {
240 panic!("boom");
241 })
242 }));
243 assert!(result.is_err());
244 assert!(get_thread_local_task_context().is_none());
245 }
246
247 #[test]
248 fn task_local_roundtrip() {
249 let rt = tokio::runtime::Builder::new_current_thread()
250 .enable_all()
251 .build()
252 .unwrap();
253 rt.block_on(async {
254 assert!(get_task_context().is_none());
255
256 let ctx = make_task_ctx();
257 let inner = with_task_context(ctx, async {
258 let c = get_task_context().unwrap();
259 assert_eq!(&*c.task_id, "task-a");
260 c
261 })
262 .await;
263
264 assert_eq!(&*inner.workflow_id, "wf-1");
265 });
266 }
267
268 #[test]
269 fn task_local_falls_back_to_thread_local() {
270 let rt = tokio::runtime::Builder::new_current_thread()
271 .enable_all()
272 .build()
273 .unwrap();
274 rt.block_on(async {
275 let ctx = make_task_ctx();
277 let result = with_thread_local_task_context(ctx, get_task_context);
278 assert!(result.is_some());
279 assert_eq!(&*result.unwrap().instance_id, "inst-1");
280 });
281 }
282
283 #[test]
284 fn macro_works() {
285 let ctx = make_task_ctx();
286 with_thread_local_task_context(ctx, || {
287 let c = task_context!().unwrap();
288 assert_eq!(&*c.task_id, "task-a");
289 });
290 }
291}