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
15pub 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 #[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 #[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 #[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 #[inline]
63 pub fn request_timeout(mut self, timeout: Duration) -> Self {
64 self.request_timeout = timeout;
65 self
66 }
67
68 #[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 #[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 #[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 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 #[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 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 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 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
186pub type BetterStackDrain<F = AcceptAll> = CloudDrain<BetterStackSerializer<F>, Bytes>;
188
189pub type BetterStackDrainHandle = CloudDrainHandle<Bytes>;
191
192pub struct BetterStackSerializer<F = AcceptAll> {
194 field_filter: InternalKVFilter<F>,
195 fallback_field: Key,
196}
197
198impl<F> BetterStackSerializer<F> {
199 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 ×tamp.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
253struct InternalKVFilter<F> {
255 inner: F,
256}
257
258impl<F> InternalKVFilter<F> {
259 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
274struct InternalClient {
276 inner: HttpClient,
277 debug: bool,
278}
279
280impl InternalClient {
281 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}