Crate future_local_storage

source ·
Expand description

Overview

This crate provides an FutureOnceCell cell-like type, which provides the similar API as the tokio::task_local but without using any macros.

Future local storage associates a value to the context of a given future. After the future finished it returns this value back to the caller. That meaning that the values is passed through the context of the executed future. This functionality can be useful for tracing async code or adding metrics to it.

Usage

use std::cell::Cell;

use future_local_storage::FutureOnceCell;

static VALUE: FutureOnceCell<Cell<u64>> = FutureOnceCell::new();

#[tokio::main]
async fn main() {
    let (output, answer) = VALUE.scope(Cell::from(0), async {
        VALUE.with(|x| {
            let value = x.get();
            x.set(value + 1);
        });

        "42".to_owned()
    }).await;

    assert_eq!(output.into_inner(), 1);
    assert_eq!(answer, "42");
}

Examples

Tracing spans

use std::{
    cell::RefCell,
    fmt::Display,
    future::Future,
    time::{Duration, Instant},
};

use future_local_storage::{FutureLocalStorage, FutureOnceCell};

/// Tracing span context.
static CONTEXT: FutureOnceCell<TracerContext> = FutureOnceCell::new();

#[derive(Debug)]
pub struct TracerContext {
    begin: Instant,
    traces: RefCell<Vec<TraceEntry>>,
}

impl TracerContext {
    /// Adds an "entered" event to the trace.
    pub fn on_enter(message: impl Display) {
        Self::with(|tracer| tracer.add_entry("entered", message));
    }

    /// Adds an "exit" event to the trace.
    pub fn on_exit(message: impl Display) {
        Self::with(|tracer| tracer.add_entry("exited", message));
    }

    /// Each future has its own tracing context in which it is executed, and after execution
    /// is complete, all events are returned along with the future output.
    pub async fn in_scope<R, F>(future: F) -> (Vec<TraceEntry>, R)
    where
        F: Future<Output = R>,
    {
        let (this, result) = future.with_scope(&CONTEXT, Self::new()).await;
        (this.traces.take(), result)
    }

    fn new() -> Self {
        let tracer = Self {
            begin: Instant::now(),
            traces: RefCell::new(Vec::new()),
        };
        tracer.add_entry("created", "a new async tracer started");
        tracer
    }

    fn with<R, F: FnOnce(&Self) -> R>(scope: F) -> R {
        CONTEXT.with(|tracer| scope(tracer))
    }

    fn add_entry(&self, event: impl Display, message: impl Display) {
        self.traces.borrow_mut().push(TraceEntry {
            duration_since: self.elapsed(),
            event: event.to_string(),
            message: message.to_string(),
        });
    }

    fn elapsed(&self) -> Duration {
        Instant::now().duration_since(self.begin)
    }
}

#[derive(Debug)]
pub struct TraceEntry {
    pub duration_since: Duration,
    pub event: String,
    pub message: String,
}

// Usage example

async fn some_method(mut a: u64) -> u64 {
    TracerContext::on_enter(format!("`some_method` with params: a={a}"));
    
    // Some async computation
     
    TracerContext::on_exit("`some_method`");
    a * 32
}

#[tokio::main]
async fn main() {
    let (trace, result) = TracerContext::in_scope(some_method(45)).await;

    println!("answer: {result}");
    println!("trace: {trace:#?}");
}

Modules

Structs

Traits