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_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}