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

use crate::log_ingestor::Log;
use crate::log_ingestor::LogIngestor;
use crate::visitor::JsonVisitor;
use serde_json::json;
use serde_json::Map;
use serde_json::Value;
use tokio::sync::mpsc::unbounded_channel;
use tracing::span;
use tracing::Subscriber;
use tracing_subscriber::registry::LookupSpan;
use tracing_subscriber::Layer;

#[allow(clippy::module_name_repetitions)]
#[derive(Debug)]
pub struct LogLayer {
    tx: Option<tokio::sync::mpsc::UnboundedSender<Log>>,
    handle: Option<std::thread::JoinHandle<()>>,
}

impl LogLayer {
    pub fn new<I>(mut ingestor: I) -> Self
    where
        I: LogIngestor + 'static,
    {
        let (tx, mut rx) = unbounded_channel::<Log>();
        // create a separate thread to manage log ingestion
        let handle = std::thread::Builder::new()
            .name(ingestor.name().into())
            .spawn(move || {
                let rt = match tokio::runtime::Builder::new_current_thread()
                    .enable_all()
                    .build()
                {
                    Err(e) => {
                        eprintln!("Runtime creation failure: {:?}", e);
                        return;
                    }
                    Ok(r) => r,
                };

                rt.block_on(async move {
                    ingestor.start();
                    while let Some(log) = rx.recv().await {
                        ingestor.ingest(log).await;
                    }
                    ingestor.flush().await;
                });
                drop(rt);
            })
            .expect("Something went wrong spawning the thread");

        Self {
            tx: Some(tx),
            handle: Some(handle),
        }
    }

    fn create_log<S: Subscriber + for<'a> LookupSpan<'a>>(
        event: &tracing::Event<'_>,
        ctx: &tracing_subscriber::layer::Context<'_, S>,
    ) -> Map<String, Value> {
        let mut log: Map<String, Value> = Map::new();
        let mut spans: Vec<Map<String, Value>> = vec![];

        if let Some(scope) = ctx.event_scope(event) {
            for span in scope.from_root() {
                let mut new_span: Map<String, Value> = Map::new();
                new_span.insert("name".to_string(), json!(span.name()));
                if let Some(fields) = span.extensions_mut().get_mut::<Map<String, Value>>() {
                    new_span.append(fields);
                }
                spans.push(new_span);
            }
        }

        // if no last span, it means there are no spans at all
        if let Some(last) = spans.last() {
            log.insert("span".to_string(), json!(last));
            log.insert("spans".to_string(), json!(spans));
        }

        log.insert(
            "level".to_string(),
            json!(event.metadata().level().as_str()),
        );
        log.insert("target".to_string(), json!(event.metadata().target()));

        if let Some(file) = event.metadata().file() {
            log.insert("file".to_string(), json!(file));
        }
        if let Some(line) = event.metadata().line() {
            log.insert("line".to_string(), json!(line));
        }

        let mut visitor = JsonVisitor::default();
        event.record(&mut visitor);

        visitor.fields.iter().for_each(|(k, v)| {
            log.insert(k.clone(), v.clone());
        });

        log.insert(
            "timestamp".to_string(),
            json!(chrono::Utc::now().to_rfc3339()),
        );

        log
    }
}

impl Drop for LogLayer {
    fn drop(&mut self) {
        // closing the channel
        if let Some(tx) = self.tx.take() {
            drop(tx);
        }
        // waiting for the thread to finish
        if let Some(handle) = self.handle.take() {
            let _result = handle.join();
        }
    }
}

impl<S> Layer<S> for LogLayer
where
    S: Subscriber + for<'a> LookupSpan<'a>,
{
    fn on_new_span(
        &self,
        attrs: &span::Attributes<'_>,
        id: &span::Id,
        ctx: tracing_subscriber::layer::Context<'_, S>,
    ) {
        let span = ctx.span(id).expect("Span not found, this is a bug");
        let mut extensions = span.extensions_mut();
        // visit values and insert them into extensions as serde_json::Map<String, serde_json::Value>
        // this way, we will be able to access them later
        let mut visitor = JsonVisitor::default();
        attrs.record(&mut visitor);
        extensions.insert(visitor.fields);
    }

    fn on_event(&self, event: &tracing::Event<'_>, ctx: tracing_subscriber::layer::Context<'_, S>) {
        // send to the channel
        if let Some(tx) = &self.tx {
            let log = Self::create_log(event, &ctx);
            if let Err(e) = tx.send(log) {
                eprintln!("LAYER: Error sending log to ingestor: {:?}", e);
            }
        }
    }
}