How to append key-value pairs to LocalKey?

[tisonkun]

In this PR I tried to implement a diagonistic fetching context info from tokio's task local values.

For thread-local, I can write:

thread_local! {
    static CONTEXT: RefCell<BTreeMap<String, String>> = const { RefCell::new(BTreeMap::new()) };
}

pub struct ThreadLocalDiagnostic {}

impl ThreadLocalDiagnostic {
    pub fn insert<K, V>(key: K, value: V)
    where
        K: Into<String>,
        V: Into<String>,
    {
        CONTEXT.with(|map| {
            map.borrow_mut().insert(key.into(), value.into());
        });
    }
}

impl Diagnostic for ThreadLocalDiagnostic {
    fn visit(&self, visitor: &mut dyn Visitor) -> Result<(), Error> {
        CONTEXT.with(|map| {
            let map = map.borrow();
            for (key, value) in map.iter() {
                let key = Key::new_ref(key.as_str());
                let value = Value::from(value);
                visitor.visit(key, value)?;
            }
            Ok(())
        })
    }
}

However, for task-local, while I write:

task_local! {
    static CONTEXT: Vec<(String, String)>;
}

pub struct TaskLocalDiagnostic {}

impl Diagnostic for TaskLocalDiagnostic {
    fn visit(&self, visitor: &mut dyn Visitor) -> Result<(), Error> {
        CONTEXT.with(|map| {
            for (k, v) in map {
                let key = Key::new_ref(k.as_str());
                let value = Value::from(v);
                visitor.visit(key, value)?;
            }
            Ok(())
        })
    }
}

pub trait FutureExt: Future {
    fn with_task_local_context(
        self,
        kvs: impl IntoIterator<Item = (String, String)>,
    ) -> impl Future<Output = Self::Output>
    where
        Self: Sized,
        Self::Output: 'static,
    {
        let mut context = CONTEXT.try_with(|v| v.clone()).unwrap_or_default();
        context.extend(kvs);
        CONTEXT.scope(context, self)
    }
}

And use it as:

    rt.block_on(
        async {
            let diag = TaskLocalDiagnostic::default();
            diag.visit(&mut PrintVisitor).unwrap();
        }
        .with_task_local_context([("user_id".to_string(), "42".to_string())])
        .with_task_local_context([("request_id".to_string(), "abc123".to_string())]),
    );

It seems the most inner future would fetch only the most inner CONTEXT and print only user_id: 42.

What I expect is:

user_id: 42
request_id: abc123

Any suggestion?

[Darksonn]

The problem is that with_task_local_context constructs the new value for the scope when constructing the future, but you need to construct it during poll of the future returned by with_task_local_context instead.

[tisonkun]

Yeah. I noticed the internal logic now. But I cannot find a way to implement what I need with a user-friendly API. Do you have some suggestions here?

That is, how the user can pass diags so that the constructed future know how to construct the task local map.

[tisonkun]

Try with:

Impl

task_local! {
    static CONTEXT: RefCell<BTreeMap<String, String>>;
}

pub struct TaskLocalDiagnostic {}

impl Diagnostic for TaskLocalDiagnostic {
    fn visit(&self, visitor: &mut dyn Visitor) -> Result<(), Error> {
        CONTEXT.with(|map| {
            let map = map.borrow();
            for (k, v) in map.iter() {
                let key = Key::new_ref(k.as_str());
                let value = Value::from(v);
                visitor.visit(key, value)?;
            }
            Ok(())
        })
    }
}

/// blablabla
pub trait FutureExt: Future + 'static {
    /// Run a future with a task-local context.
    fn with_task_local_context(
        self,
        kvs: impl IntoIterator<Item = (String, String)>,
    ) -> Pin<Box<dyn Future<Output = Self::Output>>>
    where
        Self: Sized,
        Self::Output: 'static,
    {
        match CONTEXT.try_with(|v| v.clone()) {
            Ok(v) => {
                let mut map = v.borrow_mut();
                map.extend(kvs);
                dbg!(&map);
                Box::pin(self)
            }
            Err(_) => {
                let map = RefCell::new(kvs.into_iter().collect());
                dbg!(&map);
                Box::pin(CONTEXT.scope(map, self))
            }
        }
    }
}

impl<F: Future + 'static> FutureExt for F {}

... and:

    rt.block_on(
        async {
            async {
                async {
                    let diag = TaskLocalDiagnostic::default();
                    diag.visit(&mut PrintVisitor).unwrap();
                }
                .with_task_local_context([("request_id".to_string(), "abc123".to_string())])
                .await
            }
            .with_task_local_context([("user_id".to_string(), "42".to_string())])
            .await
        }
        .with_task_local_context([("members".to_string(), "tison".to_string())]),
    );

Which output as:

[diagnostics/tokio-task-local/src/lib.rs:70:17] &map = RefCell {
    value: {
        "members": "tison",
    },
}
[diagnostics/tokio-task-local/src/lib.rs:65:17] &map = {
    "members": "tison",
    "user_id": "42",
}
[diagnostics/tokio-task-local/src/lib.rs:65:17] &map = {
    "members": "tison",
    "request_id": "abc123",
}
members: tison

[tisonkun]

Resolved in fast/logforth@c66b21e

The original requirement is task local context and I learn from tokio's tech without directly depend on tokio's task_local.