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
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
//! A mapped diagnostic context (MDC) for use with the `log` crate.
//!
//! An MDC is a thread local map of strings used to make relevant information
//! from a system available in its log messages. Logging crates such as
//! [log4rs][log4rs] will retrieve values from the MDC for output.
//!
//! For example, a web server may process many requests simultaneously on
//! different threads. Generating an ID for each request and storing it in the
//! MDC makes it easy to partition log messages on a per-request basis.
//!
//! # Examples
//!
//! Forwarding the contents of the MDC to a new thread:
//!
//! ```
//! use std::thread;
//!
//! let mut mdc = vec![];
//! log_mdc::iter(|k, v| mdc.push((k.to_owned(), v.to_owned())));
//!
//! thread::spawn(|| {
//!     log_mdc::extend(mdc);
//! });
//! ```
//!
//! [log4rs]: https://crates.io/crates/log4rs
#![doc(html_root_url="https://sfackler.github.io/rust-log-mdc/doc/v0.1.0")]
#![warn(missing_docs)]

use std::borrow::Borrow;
use std::cell::RefCell;
use std::collections::HashMap;
use std::hash::Hash;

thread_local!(static MDC: RefCell<HashMap<String, String>> = RefCell::new(HashMap::new()));

/// Inserts a new entry into the MDC, returning the old value.
pub fn insert<K, V>(key: K, value: V) -> Option<String>
    where K: Into<String>,
          V: Into<String>
{
    MDC.with(|m| m.borrow_mut().insert(key.into(), value.into()))
}

/// Inserts a new entry into the MDC in a scoped fashion.
///
/// When the returned guard falls out of scope, it will restore the old value corresponding to the
/// key.
///
/// # Examples
///
/// ```
/// let guard = log_mdc::insert_scoped("foo", "a");
/// log_mdc::get("foo", |v| assert_eq!(Some("a"), v));
///
/// drop(guard);
/// log_mdc::get("foo", |v| assert_eq!(None, v));
/// ```
///
/// ```
/// log_mdc::insert("foo", "a");
///
/// let guard = log_mdc::insert_scoped("foo", "b");
/// log_mdc::get("foo", |v| assert_eq!(Some("b"), v));
///
/// drop(guard);
/// log_mdc::get("foo", |v| assert_eq!(Some("a"), v));
/// ```
pub fn insert_scoped<K, V>(key: K, value: V) -> InsertGuard
    where K: Into<String>,
          V: Into<String>
{
    let key = key.into();
    let old_value = insert(&*key, value);

    InsertGuard {
        key: Some(key),
        old_value: old_value,
    }
}

/// Extends the MDC with new entries.
pub fn extend<K, V, I>(entries: I)
    where K: Into<String>,
          V: Into<String>,
          I: IntoIterator<Item = (K, V)>
{
    MDC.with(|m| m.borrow_mut().extend(entries.into_iter().map(|(k, v)| (k.into(), v.into()))));
}

/// Extends the MDC with new entries in a scoped fashion.
///
/// When the returned guard falls out of scope, it will restore the old values corresponding to the
/// keys.
///
/// # Examples
///
/// ```
/// log_mdc::insert("foo", "a");
///
/// let entries = [
///     ("foo", "b"),
///     ("fizz", "buzz"),
/// ];
///
/// let guard = log_mdc::extend_scoped(entries.iter().cloned());
/// log_mdc::get("foo", |v| assert_eq!(Some("b"), v));
/// log_mdc::get("fizz", |v| assert_eq!(Some("buzz"), v));
///
/// drop(guard);
/// log_mdc::get("foo", |v| assert_eq!(Some("a"), v));
/// log_mdc::get("fizz", |v| assert_eq!(None, v));
/// ```
pub fn extend_scoped<K, V, I>(entries: I) -> ExtendGuard
    where K: Into<String>,
          V: Into<String>,
          I: IntoIterator<Item = (K, V)>
{
    MDC.with(|m| {
        let mut m = m.borrow_mut();

        let old_entries = entries.into_iter()
            .map(|(k, v)| (k.into(), v.into()))
            .map(|(k, v)| {
                let v = m.insert(k.clone(), v);
                (k, v)
            })
            .collect();

        ExtendGuard(old_entries)
    })
}

/// Retrieves a value from the MDC.
pub fn get<Q: ?Sized, F, T>(key: &Q, f: F) -> T
    where String: Borrow<Q>,
          Q: Hash + Eq,
          F: FnOnce(Option<&str>) -> T
{
    MDC.with(|m| f(m.borrow().get(key).map(|v| &**v)))
}

/// Removes a value from the MDC.
pub fn remove<Q: ?Sized>(key: &Q) -> Option<String>
    where String: Borrow<Q>,
          Q: Hash + Eq
{
    MDC.with(|m| m.borrow_mut().remove(key))
}

/// Removes all values from the MDC.
pub fn clear() {
    MDC.with(|m| m.borrow_mut().clear())
}

/// Invokes the provided closure for each entry in the MDC.
pub fn iter<F>(mut f: F)
    where F: FnMut(&str, &str)
{
    MDC.with(|m| {
        for (key, value) in m.borrow().iter() {
            f(key, value)
        }
    })
}

/// A guard object which restores an MDC entry when dropped.
pub struct InsertGuard {
    key: Option<String>,
    old_value: Option<String>,
}

impl Drop for InsertGuard {
    fn drop(&mut self) {
        let key = self.key.take().unwrap();
        match self.old_value.take() {
            Some(value) => insert(key, value),
            None => remove(&key),
        };
    }
}

/// A guard objects which restores MDC entries when dropped.
pub struct ExtendGuard(Vec<(String, Option<String>)>);

impl Drop for ExtendGuard {
    fn drop(&mut self) {
        MDC.with(|m| {
            let mut m = m.borrow_mut();

            for (key, value) in self.0.drain(..) {
                match value {
                    Some(value) => m.insert(key, value),
                    None => m.remove(&key),
                };
            }
        })
    }
}