use std::future::Future;
use tokio::task_local;
use uuid::Uuid;
task_local! {
static TRACE_ID: String;
}
pub async fn with_correlation_id<F, T>(trace_id: Option<String>, fut: F) -> T
where
F: Future<Output = T>,
{
let id = trace_id.unwrap_or_else(|| Uuid::new_v4().to_string());
TRACE_ID.scope(id, fut).await
}
pub fn correlation_id() -> Option<String> {
TRACE_ID.try_with(|id| id.clone()).ok()
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn binds_a_generated_trace_id_for_the_call_tree() {
with_correlation_id(None, async {
assert!(correlation_id().is_some());
})
.await;
}
#[tokio::test]
async fn reuses_a_provided_trace_id() {
with_correlation_id(Some("fixed-id".to_string()), async {
assert_eq!(correlation_id().as_deref(), Some("fixed-id"));
})
.await;
}
#[tokio::test]
async fn reports_none_outside_any_bound_scope() {
assert_eq!(correlation_id(), None);
}
}