Skip to main content

slog_cloud/
bstack.rs

1use std::time::Duration;
2
3use bytes::Bytes;
4use chrono::{SecondsFormat, Utc};
5use slog::{KV, Key, Level, OwnedKVList, Record, Serializer};
6
7use crate::{
8    batch::NDJSONBatchBuilder,
9    client::{Client, HttpClient, Method, SendError},
10    drain::{CloudDrain, CloudDrainBuilder, CloudDrainHandle, CloudDrainTask},
11    error::Error,
12    serializer::{AcceptAll, JsonMessageBuilder, KVFilter, LogMessageSerializer},
13};
14
15/// Builder for the Better Stack log drain.
16pub struct BetterStackDrainBuilder<F = AcceptAll> {
17    field_filter: F,
18    fallback_field: Key,
19    inner: CloudDrainBuilder,
20    request_timeout: Duration,
21    debug: bool,
22}
23
24impl BetterStackDrainBuilder {
25    /// Create a new drain builder.
26    #[inline]
27    pub const fn new() -> Self {
28        Self {
29            field_filter: AcceptAll,
30            fallback_field: "",
31            inner: CloudDrainBuilder::new(),
32            request_timeout: Duration::from_secs(60),
33            debug: false,
34        }
35    }
36}
37
38impl Default for BetterStackDrainBuilder {
39    #[inline]
40    fn default() -> Self {
41        Self::new()
42    }
43}
44
45impl<F> BetterStackDrainBuilder<F> {
46    /// Set the log message queue capacity or `None` for unlimited (the default
47    /// is `None`).
48    #[inline]
49    pub fn queue_capacity(mut self, queue_capacity: Option<usize>) -> Self {
50        self.inner = self.inner.queue_capacity(queue_capacity);
51        self
52    }
53
54    /// Set the maximum number of concurrent message sends (the default is 16).
55    #[inline]
56    pub fn max_concurrency(mut self, max_concurrency: usize) -> Self {
57        self.inner = self.inner.max_concurrency(max_concurrency);
58        self
59    }
60
61    /// Set the request timeout (the default is 60 seconds).
62    #[inline]
63    pub fn request_timeout(mut self, timeout: Duration) -> Self {
64        self.request_timeout = timeout;
65        self
66    }
67
68    /// Set the initial retry delay (the default is 1 second).
69    #[inline]
70    pub fn initial_retry_delay(mut self, initial_retry_delay: Duration) -> Self {
71        self.inner = self.inner.initial_retry_delay(initial_retry_delay);
72        self
73    }
74
75    /// Set the maximum retry delay (the default is 60 seconds).
76    #[inline]
77    pub fn max_retry_delay(mut self, max_retry_delay: Duration) -> Self {
78        self.inner = self.inner.max_retry_delay(max_retry_delay);
79        self
80    }
81
82    /// Set the maximum number of retries or `None` for unlimited (the default
83    /// is unlimited).
84    #[inline]
85    pub fn max_retries(mut self, max_retries: Option<u32>) -> Self {
86        self.inner = self.inner.max_retries(max_retries);
87        self
88    }
89
90    /// Use a given key-value pair filter.
91    ///
92    /// All key-value pairs rejected by the filter will be serialized under a
93    /// given fallback field.
94    ///
95    /// This feature can be used if you want Better Stack to index only a given
96    /// subset of fields.
97    pub fn kv_filter<K, T>(self, fallback_field: K, filter: T) -> BetterStackDrainBuilder<T>
98    where
99        K: Into<Key>,
100    {
101        let mut fallback_field = fallback_field.into();
102
103        if fallback_field.is_empty() {
104            fallback_field = "misc";
105        }
106
107        BetterStackDrainBuilder {
108            field_filter: filter,
109            fallback_field,
110            inner: self.inner,
111            request_timeout: self.request_timeout,
112            debug: self.debug,
113        }
114    }
115
116    /// Enable or disable debug mode (the default is disabled).
117    ///
118    /// In the debug mode you'll be able to see some runtime info on stderr
119    /// that will help you with setting up the drain (e.g. failed requests).
120    /// With debug mode disabled, all errors will be silently ignored.
121    #[inline]
122    pub fn debug_mode(mut self, enable: bool) -> Self {
123        self.inner = self.inner.debug(enable);
124        self.debug = enable;
125        self
126    }
127
128    /// Build the drain.
129    pub fn build(
130        self,
131        ingesting_host: &str,
132        source_token: &str,
133    ) -> Result<(BetterStackDrain<F>, CloudDrainTask, BetterStackDrainHandle), Error> {
134        let client = InternalClient::new(
135            ingesting_host,
136            source_token,
137            self.request_timeout,
138            self.debug,
139        )?;
140
141        let serializer = BetterStackSerializer::new(self.field_filter, self.fallback_field);
142        let batch_builder = NDJSONBatchBuilder::new(10_000_000);
143
144        let (drain, task, handle) = self.inner.build(client, serializer, batch_builder);
145
146        Ok((drain, task, handle))
147    }
148
149    /// Build the drain and spawn a tokio task responsible for sending log
150    /// messages.
151    pub fn spawn_task(
152        self,
153        ingesting_host: &str,
154        source_token: &str,
155    ) -> Result<(BetterStackDrain<F>, BetterStackDrainHandle), Error> {
156        let (drain, task, handle) = self.build(ingesting_host, source_token)?;
157
158        tokio::spawn(task);
159
160        Ok((drain, handle))
161    }
162
163    /// Build the drain and spawn a thread responsible for sending log
164    /// messages.
165    pub fn spawn_thread(
166        self,
167        ingesting_host: &str,
168        source_token: &str,
169    ) -> Result<(BetterStackDrain<F>, BetterStackDrainHandle), Error> {
170        let (drain, task, handle) = self.build(ingesting_host, source_token)?;
171
172        std::thread::spawn(move || {
173            let runtime = tokio::runtime::Builder::new_current_thread()
174                .enable_io()
175                .enable_time()
176                .build()
177                .expect("unable to create tokio runtime");
178
179            runtime.block_on(task)
180        });
181
182        Ok((drain, handle))
183    }
184}
185
186/// Better Stack log drain.
187pub type BetterStackDrain<F = AcceptAll> = CloudDrain<BetterStackSerializer<F>, Bytes>;
188
189/// Better Stack drain handle.
190pub type BetterStackDrainHandle = CloudDrainHandle<Bytes>;
191
192/// Better Stack log message serializer.
193pub struct BetterStackSerializer<F = AcceptAll> {
194    field_filter: InternalKVFilter<F>,
195    fallback_field: Key,
196}
197
198impl<F> BetterStackSerializer<F> {
199    /// Create a new serializer with a given field filter and a fallback field.
200    fn new(field_filter: F, fallback_field: Key) -> Self {
201        Self {
202            field_filter: InternalKVFilter::new(field_filter),
203            fallback_field,
204        }
205    }
206}
207
208impl<F> LogMessageSerializer for BetterStackSerializer<F>
209where
210    F: KVFilter,
211{
212    type Serialized = Bytes;
213
214    fn serialize(
215        &self,
216        record: &Record,
217        logger_values: &OwnedKVList,
218    ) -> slog::Result<Self::Serialized> {
219        let mut builder =
220            JsonMessageBuilder::new().with_field_filter(self.fallback_field, &self.field_filter);
221
222        let level = match record.level() {
223            Level::Critical => "critical",
224            Level::Error => "error",
225            Level::Warning => "warn",
226            Level::Info => "info",
227            Level::Debug => "debug",
228            Level::Trace => "trace",
229        };
230
231        let file = record.file();
232        let line = record.line();
233
234        builder.emit_str("level", level)?;
235        builder.emit_arguments("file", &format_args!("{}:{}", file, line))?;
236        builder.emit_arguments("message", record.msg())?;
237
238        let timestamp = Utc::now();
239
240        builder.emit_str(
241            "dt",
242            &timestamp.to_rfc3339_opts(SecondsFormat::Micros, true),
243        )?;
244
245        record.kv().serialize(record, &mut builder)?;
246
247        logger_values.serialize(record, &mut builder)?;
248
249        builder.finish()
250    }
251}
252
253/// Internal key-value filter that accepts required fields.
254struct InternalKVFilter<F> {
255    inner: F,
256}
257
258impl<F> InternalKVFilter<F> {
259    /// Create a new internal key-value filter.
260    fn new(inner: F) -> Self {
261        Self { inner }
262    }
263}
264
265impl<F> KVFilter for InternalKVFilter<F>
266where
267    F: KVFilter,
268{
269    fn is_accepted(&self, key: &Key) -> bool {
270        matches!(*key, "level" | "file" | "message" | "dt") || self.inner.is_accepted(key)
271    }
272}
273
274/// Internal client for sending log messages to Better Stack.
275struct InternalClient {
276    inner: HttpClient,
277    debug: bool,
278}
279
280impl InternalClient {
281    /// Create a new client.
282    fn new(
283        ingesting_host: &str,
284        source_token: &str,
285        request_timeout: Duration,
286        debug: bool,
287    ) -> Result<Self, Error> {
288        let url = format!("https://{ingesting_host}").parse().map_err(|err| {
289            Error::from_static_msg_and_cause("unable to construct Better Stack URL", err)
290        })?;
291
292        let inner = HttpClient::builder()
293            .method(Method::POST)
294            .header("Content-Type", "application/x-ndjson")?
295            .header("Authorization", format!("Bearer {source_token}"))?
296            .request_timeout(request_timeout)
297            .build(url)?;
298
299        let res = Self { inner, debug };
300
301        Ok(res)
302    }
303}
304
305impl Client for InternalClient {
306    type Message = Bytes;
307
308    async fn send(&self, msg: Self::Message) -> Result<(), SendError> {
309        self.inner.send(msg).await.map_err(|err| {
310            if self.debug {
311                eprintln!("Better Stack request failed: {err:#}");
312            }
313
314            err.into()
315        })
316    }
317}
318
319#[cfg(test)]
320mod tests {
321    use chrono::{DateTime, Utc};
322    use slog::{Key, info, o, warn};
323
324    use crate::serializer::{AcceptAll, test_utils::serialize_logs};
325
326    use super::BetterStackSerializer;
327
328    #[test]
329    fn serializes_log_record() {
330        let serializer = BetterStackSerializer::new(AcceptAll, "misc");
331
332        let logs = serialize_logs(serializer, |log| {
333            warn!(log, "hello {}", "world"; "key" => "value");
334        });
335
336        let msg = &logs[0];
337
338        assert_eq!(msg["level"], "warn");
339        assert_eq!(msg["message"], "hello world");
340        assert_eq!(msg["key"], "value");
341
342        let file = msg["file"].as_str().unwrap();
343
344        assert!(file.starts_with(concat!(file!(), ":")), "{file}");
345
346        let timestamp: DateTime<Utc> = msg["dt"].as_str().unwrap().parse().unwrap();
347
348        assert!((Utc::now() - timestamp).num_seconds().abs() < 60);
349    }
350
351    #[test]
352    fn later_key_values_take_precedence() {
353        let serializer = BetterStackSerializer::new(AcceptAll, "misc");
354
355        let logs = serialize_logs(serializer, |log| {
356            let parent = log.new(o!("a" => "parent", "b" => "parent"));
357            let child = parent.new(o!("a" => "child"));
358
359            info!(child, "msg"; "b" => "record", "c" => "first", "c" => "second");
360        });
361
362        let msg = &logs[0];
363
364        assert_eq!(msg["a"], "child");
365        assert_eq!(msg["b"], "record");
366        assert_eq!(msg["c"], "second");
367    }
368
369    #[test]
370    fn reserved_fields_cannot_be_overridden() {
371        let serializer = BetterStackSerializer::new(AcceptAll, "misc");
372
373        let logs = serialize_logs(serializer, |log| {
374            let log = log.new(o!("level" => "fake"));
375
376            info!(log, "msg"; "message" => "fake", "file" => "fake", "dt" => "fake");
377        });
378
379        let msg = &logs[0];
380
381        assert_eq!(msg["level"], "info");
382        assert_eq!(msg["message"], "msg");
383        assert_ne!(msg["file"], "fake");
384        assert_ne!(msg["dt"], "fake");
385    }
386
387    #[test]
388    fn kv_filter_does_not_apply_to_reserved_fields() {
389        let serializer = BetterStackSerializer::new(|_: &Key| false, "extra");
390
391        let logs = serialize_logs(serializer, |log| {
392            info!(log, "msg"; "key" => "value");
393        });
394
395        let msg = &logs[0];
396
397        assert_eq!(msg["level"], "info");
398        assert_eq!(msg["message"], "msg");
399        assert!(msg["file"].is_string());
400        assert!(msg["dt"].is_string());
401        assert!(msg.get("key").is_none());
402        assert_eq!(msg["extra"], "key: value");
403    }
404}