tracing_cloudwatch/
layer.rs

1use std::sync::Arc;
2
3use tracing_core::{span, Event, Subscriber};
4use tracing_subscriber::{
5    fmt::{self, format, MakeWriter},
6    layer::Context,
7    registry::LookupSpan,
8    Layer,
9};
10
11use crate::{
12    client::CloudWatchClient,
13    dispatch::{CloudWatchDispatcher, Dispatcher, NoopDispatcher},
14    export::ExportConfig,
15};
16
17/// An AWS Cloudwatch propagation layer.
18pub struct CloudWatchLayer<S, D, N = format::DefaultFields, E = format::Format<format::Full, ()>> {
19    fmt_layer: fmt::Layer<S, N, E, Arc<D>>,
20}
21
22/// Construct [CloudWatchLayer] to compose with tracing subscriber.
23pub fn layer<S>() -> CloudWatchLayer<S, NoopDispatcher>
24where
25    S: Subscriber + for<'span> LookupSpan<'span>,
26{
27    CloudWatchLayer::default()
28}
29
30impl<S> Default
31    for CloudWatchLayer<S, NoopDispatcher, format::DefaultFields, format::Format<format::Full, ()>>
32where
33    S: Subscriber + for<'span> LookupSpan<'span>,
34{
35    fn default() -> Self {
36        CloudWatchLayer::<S,NoopDispatcher, format::DefaultFields, format::Format<format::Full,()>>::new(Arc::new(NoopDispatcher::new()))
37    }
38}
39
40impl<S, D> CloudWatchLayer<S, D, format::DefaultFields, format::Format<format::Full, ()>>
41where
42    S: Subscriber + for<'span> LookupSpan<'span>,
43    D: Dispatcher + 'static,
44    Arc<D>: for<'writer> MakeWriter<'writer>,
45{
46    pub fn new(dispatcher: Arc<D>) -> Self {
47        Self {
48            fmt_layer: fmt::Layer::default()
49                .without_time()
50                .with_writer(dispatcher)
51                .with_ansi(false)
52                .with_level(true)
53                .with_line_number(true)
54                .with_file(true)
55                .with_target(false),
56        }
57    }
58}
59
60impl<S, D, N, L, T> CloudWatchLayer<S, D, N, format::Format<L, T>>
61where
62    N: for<'writer> fmt::FormatFields<'writer> + 'static,
63{
64    /// Configure to display line number and filename.
65    /// Default true
66    pub fn with_code_location(self, display: bool) -> Self {
67        Self {
68            fmt_layer: self.fmt_layer.with_line_number(display).with_file(display),
69        }
70    }
71
72    /// Configure to display target module.
73    /// Default false.
74    pub fn with_target(self, display: bool) -> Self {
75        Self {
76            fmt_layer: self.fmt_layer.with_target(display),
77        }
78    }
79}
80
81impl<S, D, N, E> CloudWatchLayer<S, D, N, E>
82where
83    S: Subscriber + for<'span> LookupSpan<'span>,
84    D: Dispatcher + 'static,
85    Arc<D>: for<'writer> MakeWriter<'writer>,
86{
87    /// Set client.
88    pub fn with_client<Client>(
89        self,
90        client: Client,
91        export_config: ExportConfig,
92    ) -> CloudWatchLayer<S, CloudWatchDispatcher, N, E>
93    where
94        Client: CloudWatchClient + Send + Sync + 'static,
95    {
96        CloudWatchLayer {
97            fmt_layer: self
98                .fmt_layer
99                .with_writer(Arc::new(CloudWatchDispatcher::new(client, export_config))),
100        }
101    }
102
103    /// Set the [`fmt::Layer`] provided as an argument.
104    /// You can control the log format for CloudWatch by setting a pre-configured [`fmt::Layer`]
105    /// However, the specification of the writer will be overrided.
106    pub fn with_fmt_layer<N2, E2, W>(
107        self,
108        fmt_layer: fmt::Layer<S, N2, E2, W>,
109    ) -> CloudWatchLayer<S, D, N2, E2> {
110        let writer = self.fmt_layer.writer().clone();
111        CloudWatchLayer {
112            fmt_layer: fmt_layer.with_writer(writer),
113        }
114    }
115}
116
117impl<S, D, N, E> Layer<S> for CloudWatchLayer<S, D, N, E>
118where
119    S: Subscriber + for<'span> LookupSpan<'span>,
120    D: Dispatcher + 'static,
121    Arc<D>: for<'writer> MakeWriter<'writer>,
122    N: for<'writer> format::FormatFields<'writer> + 'static,
123    E: format::FormatEvent<S, N> + 'static,
124{
125    fn on_enter(&self, id: &span::Id, ctx: Context<'_, S>) {
126        self.fmt_layer.on_enter(id, ctx)
127    }
128    fn on_event(&self, event: &Event<'_>, ctx: Context<'_, S>) {
129        self.fmt_layer.on_event(event, ctx)
130    }
131
132    fn on_register_dispatch(&self, collector: &tracing::Dispatch) {
133        self.fmt_layer.on_register_dispatch(collector)
134    }
135
136    fn on_layer(&mut self, subscriber: &mut S) {
137        let _ = subscriber;
138    }
139
140    fn enabled(&self, metadata: &tracing::Metadata<'_>, ctx: Context<'_, S>) -> bool {
141        self.fmt_layer.enabled(metadata, ctx)
142    }
143
144    fn on_new_span(&self, attrs: &span::Attributes<'_>, id: &span::Id, ctx: Context<'_, S>) {
145        self.fmt_layer.on_new_span(attrs, id, ctx)
146    }
147
148    fn on_record(&self, id: &span::Id, values: &span::Record<'_>, ctx: Context<'_, S>) {
149        self.fmt_layer.on_record(id, values, ctx)
150    }
151
152    fn on_follows_from(&self, span: &span::Id, follows: &span::Id, ctx: Context<'_, S>) {
153        self.fmt_layer.on_follows_from(span, follows, ctx)
154    }
155
156    fn event_enabled(&self, event: &Event<'_>, ctx: Context<'_, S>) -> bool {
157        self.fmt_layer.event_enabled(event, ctx)
158    }
159
160    fn on_exit(&self, id: &span::Id, ctx: Context<'_, S>) {
161        self.fmt_layer.on_exit(id, ctx)
162    }
163
164    fn on_close(&self, id: span::Id, ctx: Context<'_, S>) {
165        self.fmt_layer.on_close(id, ctx)
166    }
167
168    fn on_id_change(&self, old: &span::Id, new: &span::Id, ctx: Context<'_, S>) {
169        self.fmt_layer.on_id_change(old, new, ctx)
170    }
171}
172
173#[cfg(test)]
174mod tests {
175    use std::sync::Mutex;
176
177    use chrono::{DateTime, TimeZone, Utc};
178    use tracing_subscriber::layer::SubscriberExt;
179
180    use crate::dispatch::LogEvent;
181
182    use super::*;
183
184    struct TestDispatcher {
185        events: Mutex<Vec<LogEvent>>,
186    }
187
188    impl TestDispatcher {
189        fn new() -> Self {
190            Self {
191                events: Mutex::new(Vec::new()),
192            }
193        }
194    }
195
196    impl Dispatcher for TestDispatcher {
197        fn dispatch(&self, input: crate::dispatch::LogEvent) {
198            self.events.lock().unwrap().push(input)
199        }
200    }
201
202    impl std::io::Write for &TestDispatcher {
203        fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
204            let timestamp: DateTime<Utc> = Utc.timestamp_opt(1_5000_000_000, 0).unwrap();
205
206            let message = String::from_utf8_lossy(buf).to_string();
207
208            self.events
209                .lock()
210                .unwrap()
211                .push(LogEvent { timestamp, message });
212
213            Ok(buf.len())
214        }
215
216        fn flush(&mut self) -> std::io::Result<()> {
217            Ok(())
218        }
219    }
220
221    #[test]
222    fn format() {
223        let dispatcher = Arc::new(TestDispatcher::new());
224        let subscriber = tracing_subscriber::registry().with(
225            CloudWatchLayer::new(dispatcher.clone())
226                .with_code_location(false)
227                .with_target(false),
228        );
229
230        tracing::subscriber::with_default(subscriber, || {
231            tracing::info_span!("span-1", xxx = "yyy").in_scope(|| {
232                tracing::debug_span!("span-2", key = "value").in_scope(|| {
233                    tracing::info!("Hello!");
234                })
235            });
236
237            tracing::error!("Error");
238        });
239
240        let dispatched = dispatcher.events.lock().unwrap().remove(0);
241        assert_eq!(
242            dispatched.message,
243            " INFO span-1{xxx=\"yyy\"}:span-2{key=\"value\"}: Hello!\n"
244        );
245
246        let dispatched = dispatcher.events.lock().unwrap().remove(0);
247        assert_eq!(dispatched.message, "ERROR Error\n");
248    }
249
250    #[test]
251    fn with_fmt_layer_json() {
252        let dispatcher = Arc::new(TestDispatcher::new());
253        let subscriber = tracing_subscriber::registry().with(
254            CloudWatchLayer::new(dispatcher.clone())
255                .with_fmt_layer(fmt::layer().json().without_time()),
256        );
257
258        tracing::subscriber::with_default(subscriber, || {
259            tracing::info_span!("span-1", xxx = "yyy").in_scope(|| {
260                tracing::debug_span!("span-2", key = "value").in_scope(|| {
261                    tracing::info!("Hello!");
262                })
263            });
264        });
265
266        let dispatched = dispatcher.events.lock().unwrap().remove(0);
267        insta::assert_debug_snapshot!(dispatched.message);
268    }
269}
tracing_cloudwatch/layer.rs

tracing_cloudwatch/
layer.rs
