daemonbit_test/
lib.rs

1use std::{
2    io,
3    ops::{Deref, DerefMut},
4    sync::Arc,
5};
6
7use memchr::memchr;
8use parking_lot::Mutex;
9use tracing::Dispatch;
10use tracing_subscriber::FmtSubscriber;
11
12pub fn capture_tracing<F, T>(f: F) -> CapturedTracing<T>
13where
14    F: FnOnce() -> T,
15{
16    let buf = InnerBuffer::new();
17    let dispatch: Dispatch = FmtSubscriber::builder()
18        // .with_env_filter("trace")
19        .with_writer({
20            let buf = buf.clone();
21            move || buf.clone()
22        })
23        .with_level(true)
24        .with_ansi(false)
25        .into();
26    let result = tracing::dispatcher::with_default(&dispatch, f);
27    CapturedTracing {
28        result,
29        buf: buf.into_inner(),
30    }
31}
32
33pub struct CapturedTracing<T> {
34    result: T,
35    buf: Vec<u8>,
36}
37#[allow(unused)]
38impl<T> CapturedTracing<T> {
39    pub fn into_result(self) -> T {
40        self.result
41    }
42    pub fn logged<S: AsRef<str>>(&self, s: S) -> bool {
43        let s = s.as_ref().as_bytes();
44        let (needle, rest) = match *s {
45            [needle, ref rest @ ..] => (needle, rest),
46            _ => return false,
47        };
48
49        let mut haystack = &*self.buf;
50        while let Some(p) = memchr(needle, haystack) {
51            haystack = &haystack[p + 1..];
52            if rest.len() > haystack.len() {
53                break;
54            }
55            if &haystack[..rest.len()] == rest {
56                return true;
57            }
58        }
59        false
60    }
61}
62impl<T> Deref for CapturedTracing<T> {
63    type Target = T;
64    fn deref(&self) -> &Self::Target {
65        &self.result
66    }
67}
68impl<T> DerefMut for CapturedTracing<T> {
69    fn deref_mut(&mut self) -> &mut Self::Target {
70        &mut self.result
71    }
72}
73
74#[derive(Clone)]
75struct InnerBuffer {
76    shared: Arc<Mutex<Vec<u8>>>,
77}
78impl InnerBuffer {
79    fn new() -> Self {
80        Self {
81            shared: Arc::new(Mutex::new(Vec::new())),
82        }
83    }
84    fn into_inner(self) -> Vec<u8> {
85        let mut shared = self.shared.lock();
86        std::mem::take(&mut *shared)
87    }
88}
89impl io::Write for InnerBuffer {
90    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
91        self.shared.lock().write(buf)
92    }
93    fn flush(&mut self) -> io::Result<()> {
94        Ok(())
95    }
96}