dcontext 0.8.0

Distributed context propagation for Rust — scoped, type-safe, serializable
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152

//! Context inheritance — propagating context across task/thread boundaries.
//!
//! Three spawn helpers cover the common cross-boundary scenarios:
//!
//! | Helper | Source | Target |
//! |--------|--------|--------|
//! | [`spawn_with_async_context`] | async (task-local) | async task |
//! | [`spawn_blocking_with_async_context`] | async (task-local) | blocking thread |
//! | [`spawn_with_sync_context`] | sync (thread-local) | async task |
//!
//! Each takes a [`ContextInheritance`] mode:
//!
//! - **Fork** — cheap Arc sharing. Child sees parent's values at fork time.
//!   Reads fall through to the frozen parent; writes are isolated (COW).
//! - **Snapshot** — full copy. Walks the entire scope chain, clones all
//!   effective values into a flat HashMap. O(1) reads in the child.

use std::cell::Cell;
use std::sync::Arc;

use crate::async_ctx::{self, TASK_CONTEXT};
use crate::snapshot::ContextSnapshot;
use crate::store::ContextStore;
use crate::sync_ctx;

/// Controls how context is captured when crossing task/thread boundaries.
///
/// - **Fork** — lightweight Arc sharing. The child's scope inherits from
///   a frozen parent node. Reads fall through; writes are isolated (COW).
///   Cheapest option when the child mostly reads.
/// - **Snapshot** — full deep copy. All effective values are collected into
///   a flat HashMap. More expensive to create but O(1) reads in the child.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum ContextInheritance {
    /// Lightweight Arc sharing — O(current_scope_keys) to create.
    #[default]
    Fork,
    /// Full value copy — O(depth × keys) to create, O(1) reads.
    Snapshot,
}

// ── Spawn helpers ─────────────────────────────────────────────

/// Spawn a Tokio task that inherits the current **async** (task-local) context.
///
/// # Example
///
/// ```rust,ignore
/// use dcontext::{spawn_with_async_context, ContextInheritance};
///
/// // Fork (default, cheapest)
/// spawn_with_async_context(ContextInheritance::Fork, async {
///     let rid: String = dcontext::async_ctx::get_context("request_id").unwrap();
/// });
///
/// // Snapshot (full copy)
/// spawn_with_async_context(ContextInheritance::Snapshot, async { /* ... */ });
/// ```
pub fn spawn_with_async_context<F>(
    mode: ContextInheritance,
    future: F,
) -> tokio::task::JoinHandle<F::Output>
where
    F: std::future::Future + Send + 'static,
    F::Output: Send + 'static,
{
    let store = capture_async(mode);
    tokio::spawn(TASK_CONTEXT.scope(Cell::new(Some(store)), future))
}

/// Spawn a blocking thread that inherits the current **async** (task-local) context.
///
/// The inherited context is installed into the blocking thread's thread-local
/// store and restored when the closure returns.
///
/// # Example
///
/// ```rust,ignore
/// use dcontext::{spawn_blocking_with_async_context, ContextInheritance};
///
/// spawn_blocking_with_async_context(ContextInheritance::Fork, || {
///     let rid: String = dcontext::sync_ctx::get_context("request_id").unwrap();
/// });
/// ```
pub fn spawn_blocking_with_async_context<F, R>(
    mode: ContextInheritance,
    f: F,
) -> tokio::task::JoinHandle<R>
where
    F: FnOnce() -> R + Send + 'static,
    R: Send + 'static,
{
    let store = capture_async(mode);
    tokio::task::spawn_blocking(move || {
        // Swap into this thread's thread-local store; restore on exit.
        let saved = sync_ctx::try_apply(|s| std::mem::replace(s, store));
        let result = f();
        if let Some(original) = saved {
            sync_ctx::try_apply(|s| *s = original);
        }
        result
    })
}

/// Spawn a Tokio task that inherits the current **sync** (thread-local) context.
///
/// # Example
///
/// ```rust,ignore
/// use dcontext::{spawn_with_sync_context, ContextInheritance};
///
/// spawn_with_sync_context(ContextInheritance::Fork, async {
///     let rid: String = dcontext::async_ctx::get_context("request_id").unwrap();
/// });
/// ```
pub fn spawn_with_sync_context<F>(
    mode: ContextInheritance,
    future: F,
) -> tokio::task::JoinHandle<F::Output>
where
    F: std::future::Future + Send + 'static,
    F::Output: Send + 'static,
{
    let store = capture_sync(mode);
    tokio::spawn(TASK_CONTEXT.scope(Cell::new(Some(store)), future))
}

// ── Internal: capture context into a ContextStore ─────────────

fn capture_async(mode: ContextInheritance) -> ContextStore {
    match mode {
        ContextInheritance::Fork => async_ctx::fork().unwrap_or_else(ContextStore::new),
        ContextInheritance::Snapshot => snapshot_to_store(async_ctx::snapshot()),
    }
}

fn capture_sync(mode: ContextInheritance) -> ContextStore {
    match mode {
        ContextInheritance::Fork => sync_ctx::fork().unwrap_or_else(ContextStore::new),
        ContextInheritance::Snapshot => snapshot_to_store(sync_ctx::snapshot()),
    }
}

fn snapshot_to_store(snap: ContextSnapshot) -> ContextStore {
    let chain = snap.scope_chain.clone();
    let values = snap
        .values
        .iter()
        .map(|(k, v)| (*k, Arc::clone(v)))
        .collect();
    ContextStore::from_values_with_chain(values, chain)
}