Skip to main content

datadog_apm_sync/
client.rs

1use crate::{api::RawSpan, model::Span};
2
3use atomic_float::AtomicF64;
4use attohttpc;
5use chrono::{DateTime, Duration, Utc};
6use log::{Level as LogLevel, Log, Record};
7use serde_json::to_string;
8use std::{
9    cell::Cell,
10    collections::{HashMap, VecDeque},
11    sync::{
12        atomic::{AtomicU16, AtomicU32, Ordering},
13        mpsc,
14    },
15};
16
17#[cfg(feature = "json")]
18use log::kv;
19
20/// Configuration settings for the client.
21#[derive(Clone, Debug)]
22pub struct Config {
23    /// Datadog apm service name
24    pub service: String,
25    /// Datadog apm environment
26    pub env: Option<String>,
27    /// Datadog agent host/ip, defaults to `localhost`.
28    pub host: String,
29    /// Datadog agent port, defaults to `8126`.
30    pub port: String,
31    /// Optional Logging Config to also set this tracer as the main logger
32    pub logging_config: Option<LoggingConfig>,
33    /// APM Config to set up APM Analytics (default is to disable)
34    pub apm_config: ApmConfig,
35    /// Turn on tracing
36    pub enable_tracing: bool,
37    /// Number of threads to send the HTTP messages to the Datadog agent
38    pub num_client_send_threads: u32,
39}
40
41impl Default for Config {
42    fn default() -> Self {
43        Config {
44            env: None,
45            host: "localhost".to_string(),
46            port: "8126".to_string(),
47            service: "".to_string(),
48            logging_config: None,
49            apm_config: ApmConfig::default(),
50            enable_tracing: false,
51            num_client_send_threads: 4,
52        }
53    }
54}
55
56#[derive(Clone, Debug)]
57pub struct LoggingConfig {
58    pub level: LogLevel,
59    pub time_format: String,
60    pub mod_filter: Vec<&'static str>,
61    pub body_filter: Vec<&'static str>,
62}
63
64impl Default for LoggingConfig {
65    fn default() -> Self {
66        LoggingConfig {
67            level: LogLevel::Info,
68            time_format: "%Y-%m-%d %H:%M:%S%z".to_string(),
69            mod_filter: Vec::new(),
70            body_filter: Vec::new(),
71        }
72    }
73}
74
75#[derive(Clone, Debug)]
76pub struct ApmConfig {
77    pub apm_enabled: bool,
78    pub sample_priority: f64,
79    pub sample_rate: f64,
80}
81
82impl Default for ApmConfig {
83    fn default() -> Self {
84        ApmConfig {
85            apm_enabled: false,
86            sample_rate: 0f64,
87            sample_priority: 0f64,
88        }
89    }
90}
91
92type TimeInNanos = u64;
93type ThreadId = u32;
94type TraceId = u64;
95type SpanId = u64;
96
97#[derive(Clone, Debug)]
98struct LogRecord {
99    pub thread_id: ThreadId,
100    pub level: tracing::Level,
101    pub time: DateTime<Utc>,
102    pub msg_str: String,
103    pub module: Option<String>,
104    #[cfg(feature = "json")]
105    pub key_values: HashMap<String, String>,
106}
107
108#[derive(Clone, Debug)]
109enum TraceCommand {
110    Log(LogRecord),
111    NewSpan(TimeInNanos, NewSpanData),
112    Enter(TimeInNanos, ThreadId, SpanId),
113    Exit(TimeInNanos, SpanId),
114    CloseSpan(TimeInNanos, SpanId),
115    Event(EventRecord),
116}
117
118#[derive(Clone, Debug)]
119struct EventRecord {
120    thread_id: ThreadId,
121    fields: HashMap<String, String>,
122    time: DateTime<Utc>,
123    level: tracing::Level,
124    module: Option<String>,
125    message: Option<String>,
126    send_trace: Option<u64>,
127}
128
129impl EventRecord {
130    pub fn new(
131        thread_id: ThreadId,
132        mut fields: HashMap<String, String>,
133        time: DateTime<Utc>,
134        level: tracing::Level,
135        module: Option<String>,
136    ) -> Self {
137        let send_trace = fields
138            .remove("send_trace")
139            .and_then(|tr_str| tr_str.parse::<u64>().ok());
140        let message = fields.remove("message");
141        EventRecord {
142            thread_id,
143            fields,
144            time,
145            level,
146            module,
147            send_trace,
148            message,
149        }
150    }
151}
152
153#[derive(Debug, Clone)]
154struct NewSpanData {
155    pub trace_id: TraceId,
156    pub id: SpanId,
157    pub name: String,
158    pub resource: String,
159    pub start: DateTime<Utc>,
160}
161
162#[derive(Clone, Debug)]
163struct SpanCollection {
164    completed_spans: Vec<Span>,
165    parent_span: Span,
166    current_spans: VecDeque<Span>,
167    entered_spans: VecDeque<u64>,
168}
169
170impl SpanCollection {
171    fn new(parent_span: Span) -> Self {
172        SpanCollection {
173            completed_spans: vec![],
174            parent_span,
175            current_spans: VecDeque::new(),
176            entered_spans: VecDeque::new(),
177        }
178    }
179
180    // Open a span by inserting the span into the "current" span map by ID.
181    fn start_span(&mut self, span: Span) {
182        let parent_id = Some(self.current_span_id().unwrap_or(self.parent_span.id));
183        self.current_spans.push_back(Span { parent_id, ..span });
184    }
185
186    // Move span to "completed" based on ID.
187    fn end_span(&mut self, nanos: u64, span_id: SpanId) {
188        let pos = self.current_spans.iter().rposition(|i| i.id == span_id);
189        if let Some(i) = pos {
190            if let Some(span) = self.current_spans.remove(i) {
191                self.completed_spans.push(Span {
192                    duration: Duration::nanoseconds(
193                        nanos as i64 - span.start.timestamp_nanos_opt().unwrap_or(0),
194                    ),
195                    ..span
196                });
197            }
198        }
199    }
200
201    // Enter a span (mark it on stack)
202    fn enter_span(&mut self, span_id: SpanId) {
203        self.entered_spans.push_back(span_id);
204    }
205
206    // Exit a span (pop from stack)
207    fn exit_span(&mut self, span_id: SpanId) {
208        let pos = self.entered_spans.iter().rposition(|i| *i == span_id);
209        if let Some(i) = pos {
210            self.entered_spans.remove(i);
211        }
212    }
213
214    /// Get the id, if present, of the most current span for this trace
215    fn current_span_id(&self) -> Option<u64> {
216        self.entered_spans.back().copied()
217    }
218
219    fn add_tag(&mut self, k: String, v: String) {
220        if let Some(span) = self.current_spans.back_mut() {
221            span.tags.insert(k.clone(), v.clone());
222        }
223        self.parent_span.tags.insert(k, v);
224    }
225
226    fn drain_current(mut self) -> Self {
227        std::mem::take(&mut self.current_spans)
228            .into_iter()
229            .for_each(|span| {
230                self.completed_spans.push(Span {
231                    duration: Utc::now().signed_duration_since(span.start),
232                    ..span
233                })
234            });
235        self
236    }
237
238    fn drain(self, end_time: DateTime<Utc>) -> Vec<Span> {
239        let parent_span = Span {
240            duration: end_time.signed_duration_since(self.parent_span.start),
241            ..self.parent_span.clone()
242        };
243        let mut ret = self.drain_current().completed_spans;
244        ret.push(parent_span);
245        ret
246    }
247}
248
249struct SpanStorage {
250    traces: HashMap<TraceId, SpanCollection>,
251    spans_to_trace_id: HashMap<SpanId, TraceId>,
252    current_trace_for_thread: HashMap<ThreadId, TraceId>,
253    current_thread_for_trace: HashMap<TraceId, ThreadId>,
254}
255
256impl SpanStorage {
257    fn new() -> Self {
258        SpanStorage {
259            traces: HashMap::new(),
260            spans_to_trace_id: HashMap::new(),
261            current_trace_for_thread: HashMap::new(),
262            current_thread_for_trace: HashMap::new(),
263        }
264    }
265
266    // Either start a new trace with the span's trace ID (if there is no span already
267    // pushed for that trace ID), or push the span on the "current" stack of spans for that
268    // trace ID.  If "parent" is true, that means we need a parent span pushed for this to
269    // represent the entire trace.
270    fn start_span(&mut self, span: Span) {
271        let trace_id = span.trace_id;
272        self.spans_to_trace_id.insert(span.id, span.trace_id);
273        if let Some(ss) = self.traces.get_mut(&trace_id) {
274            ss.start_span(span);
275        } else {
276            self.traces.insert(trace_id, SpanCollection::new(span));
277        }
278    }
279
280    /// End a span and update the current "top of the stack"
281    fn end_span(&mut self, nanos: u64, span_id: SpanId) {
282        if let Some(trace_id) = self.spans_to_trace_id.remove(&span_id) {
283            if let Some(ref mut ss) = self.traces.get_mut(&trace_id) {
284                ss.end_span(nanos, span_id);
285            }
286        }
287    }
288
289    /// Enter a span for trace, and keep track so that new spans get the correct parent.
290    /// Keep track of which trace the current thread is in (for logging and events)
291    fn enter_span(&mut self, thread_id: ThreadId, span_id: SpanId) {
292        let t_id = self.spans_to_trace_id.get(&span_id).copied();
293        if let Some(trace_id) = t_id {
294            if let Some(ref mut ss) = self.traces.get_mut(&trace_id) {
295                ss.enter_span(span_id);
296                if ss.entered_spans.len() == 1 {
297                    self.set_current_trace(thread_id, trace_id);
298                }
299            }
300        }
301    }
302
303    /// Exit a span for trace, and keep track so that new spans get the correct parent
304    fn exit_span(&mut self, span_id: SpanId) {
305        let trace_id = self.spans_to_trace_id.get(&span_id).cloned();
306        if let Some(trace_id) = trace_id {
307            if let Some(ref mut ss) = self.traces.get_mut(&trace_id) {
308                ss.exit_span(span_id);
309                if ss.entered_spans.is_empty() {
310                    self.remove_current_trace(trace_id);
311                }
312            }
313        }
314    }
315
316    /// Drain the span collection for this trace so we can send the trace through to Datadog,
317    /// This effectively ends the trace.  Any new spans on this trace ID will have the same
318    /// trace ID, but have a new parent span (and a new trace line in Datadog).
319    fn drain_completed(&mut self, trace_id: TraceId, end: DateTime<Utc>) -> Vec<Span> {
320        if let Some(ss) = self.traces.remove(&trace_id) {
321            ss.drain(end)
322        } else {
323            vec![]
324        }
325    }
326
327    /// Record tag info onto a span
328    fn span_record_tag(&mut self, trace_id: TraceId, key: String, value: String) {
329        if let Some(ref mut ss) = self.traces.get_mut(&trace_id) {
330            ss.add_tag(key, value)
331        }
332    }
333
334    fn get_trace_id_for_thread(&self, thread_id: ThreadId) -> Option<u64> {
335        self.current_trace_for_thread.get(&thread_id).copied()
336    }
337
338    fn set_current_trace(&mut self, thread_id: ThreadId, trace_id: TraceId) {
339        self.current_trace_for_thread.insert(thread_id, trace_id);
340        self.current_thread_for_trace.insert(trace_id, thread_id);
341    }
342
343    fn remove_current_trace(&mut self, trace_id: TraceId) {
344        let thread_id = self.current_thread_for_trace.remove(&trace_id);
345        if let Some(thr) = thread_id {
346            self.current_trace_for_thread.remove(&thr);
347        }
348    }
349
350    /// Get the id, if present, of the most current span for the given trace
351    fn current_span_id(&self, trace_id: TraceId) -> Option<SpanId> {
352        self.traces.get(&trace_id).and_then(|s| s.current_span_id())
353    }
354}
355
356fn filter_log(storage: &SpanStorage, log_config: &LoggingConfig, record: LogRecord) {
357    let mod_skip = record
358        .module
359        .as_ref()
360        .map(|m: &String| {
361            log_config
362                .mod_filter
363                .iter()
364                .any(|filter| m.contains(*filter))
365        })
366        .unwrap_or(false);
367    let body_skip = log_config
368        .body_filter
369        .iter()
370        .any(|f| record.msg_str.contains(*f));
371    if !mod_skip && !body_skip {
372        let log_body = build_log_body(&record);
373        match storage
374            .get_trace_id_for_thread(record.thread_id)
375            .and_then(|tr_id| storage.current_span_id(tr_id).map(|sp_id| (tr_id, sp_id)))
376        {
377            Some((tr, sp)) => {
378                // Both trace and span are active on this thread
379                println!(
380                    "{time} {level} [trace-id:{traceid} span-id:{spanid}] [{module}] {body}",
381                    time = record.time.format(log_config.time_format.as_ref()),
382                    traceid = tr,
383                    spanid = sp,
384                    level = record.level,
385                    module = record.module.clone().unwrap_or("-".to_string()),
386                    body = log_body
387                );
388            }
389            _ => {
390                // Both trace and span are not active on this thread
391                println!(
392                    "{time} {level} [{module}] {body}",
393                    time = record.time.format(log_config.time_format.as_ref()),
394                    level = record.level,
395                    module = record.module.clone().unwrap_or("-".to_string()),
396                    body = log_body
397                );
398            }
399        }
400    }
401}
402
403fn trace_server_loop(
404    client: DdAgentClient,
405    buffer_receiver: mpsc::Receiver<TraceCommand>,
406    log_config: Option<LoggingConfig>,
407) {
408    let mut storage = SpanStorage::new();
409
410    loop {
411        match buffer_receiver.recv() {
412            Ok(TraceCommand::Log(record)) => {
413                log_config
414                    .as_ref()
415                    .inspect(|lc| filter_log(&storage, lc, record));
416            }
417            Ok(TraceCommand::NewSpan(_nanos, data)) => {
418                storage.start_span(Span {
419                    id: data.id,
420                    trace_id: data.trace_id,
421                    tags: HashMap::new(),
422                    parent_id: None,
423                    start: data.start,
424                    name: data.name,
425                    resource: data.resource,
426                    sql: None,
427                    duration: Duration::seconds(0),
428                });
429            }
430            Ok(TraceCommand::Enter(_nanos, thread_id, span_id)) => {
431                storage.enter_span(thread_id, span_id);
432            }
433            Ok(TraceCommand::Exit(_nanos, span_id)) => {
434                storage.exit_span(span_id);
435            }
436            Ok(TraceCommand::Event(mut event)) => {
437                // Events are only valid if the trace_id flag is set
438                // Send trace specified the trace to send, so use that instead of the thread's
439                // current trace.
440                if let Some(send_trace_id) = event.send_trace {
441                    let send_vec = storage.drain_completed(send_trace_id, event.time);
442                    // Thread has ended this trace.  Until it enters a new span, it
443                    // is not in a trace.
444                    storage.remove_current_trace(send_trace_id);
445                    if !send_vec.is_empty() {
446                        client.send(send_vec);
447                    }
448                }
449
450                // Tag events only work inside a trace, so get the trace from the thread.
451                // No trace means no tagging.
452                let trace_id_opt = storage.get_trace_id_for_thread(event.thread_id);
453                if let Some(trace_id) = trace_id_opt {
454                    if let Some(type_event) = event.fields.remove("error.etype") {
455                        storage.span_record_tag(trace_id, "error.type".to_string(), type_event)
456                    }
457                    event.fields.iter().for_each(|(k, v)| {
458                        if k != "message" {
459                            storage.span_record_tag(trace_id, k.clone(), v.clone());
460                        }
461                    });
462                }
463
464                if let Some(ref lc) = log_config {
465                    if let Some(msg_str) = event.message {
466                        let record = LogRecord {
467                            thread_id: event.thread_id,
468                            level: event.level,
469                            time: event.time,
470                            msg_str,
471                            module: event.module,
472                            #[cfg(feature = "json")]
473                            key_values: event.fields,
474                        };
475                        filter_log(&storage, lc, record);
476                    }
477                }
478            }
479            Ok(TraceCommand::CloseSpan(nanos, span_id)) => {
480                storage.end_span(nanos, span_id);
481            }
482            Err(_) => {
483                return;
484            }
485        }
486    }
487}
488
489fn build_log_body(record: &LogRecord) -> String {
490    #[cfg(not(feature = "json"))]
491    {
492        record.msg_str.clone()
493    }
494    #[cfg(feature = "json")]
495    {
496        if record.key_values.is_empty() {
497            record.msg_str.clone()
498        } else {
499            let mut body = HashMap::new();
500            body.insert("message".to_string(), record.msg_str.clone());
501            for (k, v) in &record.key_values {
502                body.insert(k.clone(), v.clone());
503            }
504            serde_json::to_string(&body).unwrap_or_else(|_| "".to_string())
505        }
506    }
507}
508
509#[derive(Debug, Clone)]
510pub struct DatadogTracing {
511    buffer_sender: mpsc::Sender<TraceCommand>,
512    log_config: Option<LoggingConfig>,
513}
514
515unsafe impl Sync for DatadogTracing {}
516
517impl DatadogTracing {
518    pub fn new(config: Config) -> DatadogTracing {
519        let (buffer_sender, buffer_receiver) = mpsc::channel();
520        let sample_rate = config.apm_config.sample_rate;
521        let client = DdAgentClient::new(&config);
522
523        let log_config = config.logging_config.clone();
524        std::thread::spawn(move || trace_server_loop(client, buffer_receiver, log_config));
525
526        let tracer = DatadogTracing {
527            buffer_sender,
528            log_config: config.logging_config,
529        };
530
531        if let Some(ref lc) = tracer.log_config {
532            let _ = log::set_boxed_logger(Box::new(tracer.clone()));
533            log::set_max_level(lc.level.to_level_filter());
534        }
535        if config.enable_tracing {
536            // Only set the global sample rate once when the tracer is set as the global tracer.
537            // This must be marked unsafe because we are overwriting a global, but it only gets done
538            // once in a process's lifetime.
539            SAMPLING_RATE.store(sample_rate, Ordering::Release);
540        }
541        tracer
542    }
543
544    pub fn get_global_sampling_rate() -> f64 {
545        SAMPLING_RATE.load(Ordering::Acquire)
546    }
547
548    fn send_log(&self, record: LogRecord) -> Result<(), ()> {
549        self.buffer_sender
550            .send(TraceCommand::Log(record))
551            .map(|_| ())
552            .map_err(|_| ())
553    }
554
555    fn send_new_span(&self, nanos: u64, span: NewSpanData) -> Result<(), ()> {
556        self.buffer_sender
557            .send(TraceCommand::NewSpan(nanos, span))
558            .map(|_| ())
559            .map_err(|_| ())
560    }
561
562    fn send_enter_span(&self, nanos: u64, thread_id: ThreadId, id: SpanId) -> Result<(), ()> {
563        self.buffer_sender
564            .send(TraceCommand::Enter(nanos, thread_id, id))
565            .map(|_| ())
566            .map_err(|_| ())
567    }
568
569    fn send_exit_span(&self, nanos: u64, id: SpanId) -> Result<(), ()> {
570        self.buffer_sender
571            .send(TraceCommand::Exit(nanos, id))
572            .map(|_| ())
573            .map_err(|_| ())
574    }
575
576    fn send_close_span(&self, nanos: u64, span_id: SpanId) -> Result<(), ()> {
577        self.buffer_sender
578            .send(TraceCommand::CloseSpan(nanos, span_id))
579            .map(|_| ())
580            .map_err(|_| ())
581    }
582
583    fn send_event(
584        &self,
585        thread_id: ThreadId,
586        event: HashMap<String, String>,
587        time: DateTime<Utc>,
588        level: &tracing::Level,
589        module: Option<String>,
590    ) -> Result<(), ()> {
591        self.buffer_sender
592            .send(TraceCommand::Event(EventRecord::new(
593                thread_id, event, time, *level, module,
594            )))
595            .map(|_| ())
596            .map_err(|_| ())
597    }
598}
599
600fn log_level_to_trace_level(level: log::Level) -> tracing::Level {
601    use log::Level::*;
602    match level {
603        Error => tracing::Level::INFO,
604        Warn => tracing::Level::INFO,
605        Info => tracing::Level::INFO,
606        Debug => tracing::Level::DEBUG,
607        Trace => tracing::Level::TRACE,
608    }
609}
610
611static UNIQUEID_COUNTER: AtomicU16 = AtomicU16::new(0);
612static THREAD_COUNTER: AtomicU32 = AtomicU32::new(0);
613
614static SAMPLING_RATE: AtomicF64 = AtomicF64::new(0.0);
615
616thread_local! {
617    static THREAD_ID: ThreadId = THREAD_COUNTER.fetch_add(1, Ordering::Relaxed);
618    static CURRENT_SPAN_ID: Cell<Option<SpanId>> = const { Cell::new(None) }
619}
620
621pub fn get_thread_id() -> ThreadId {
622    THREAD_ID.with(|id| *id)
623}
624
625pub fn get_current_span_id() -> Option<SpanId> {
626    CURRENT_SPAN_ID.with(|id| id.get())
627}
628
629pub fn set_current_span_id(new_id: Option<SpanId>) {
630    CURRENT_SPAN_ID.with(|id| {
631        id.set(new_id);
632    })
633}
634
635// Format
636// |                       6 bytes                       |      2 bytes    |
637// +--------+--------+--------+--------+--------+--------+--------+--------+
638// |     number of milliseconds since epoch (1970)       | static counter  |
639// +--------+--------+--------+--------+--------+--------+--------+--------+
640// 0        8        16       24       32       40       48       56       64
641//
642// This will hold up to the year 10,000 before it cycles.
643pub fn create_unique_id64() -> u64 {
644    let millis_since_epoch = (Utc::now().timestamp_millis() << 16) as u64;
645    millis_since_epoch + UNIQUEID_COUNTER.fetch_add(1, Ordering::Relaxed) as u64
646}
647
648pub struct HashMapVisitor {
649    fields: HashMap<String, String>,
650}
651
652impl HashMapVisitor {
653    fn new() -> Self {
654        // Event/Span vectors should never have more than ten fields.
655        HashMapVisitor {
656            fields: HashMap::new(),
657        }
658    }
659    fn add_value(&mut self, field: &tracing::field::Field, value: String) {
660        self.fields.insert(field.name().to_string(), value);
661    }
662}
663
664impl tracing::field::Visit for HashMapVisitor {
665    fn record_str(&mut self, field: &tracing::field::Field, value: &str) {
666        self.add_value(field, value.to_string());
667    }
668    fn record_bool(&mut self, field: &tracing::field::Field, value: bool) {
669        self.add_value(field, value.to_string());
670    }
671    fn record_u64(&mut self, field: &tracing::field::Field, value: u64) {
672        self.add_value(field, value.to_string());
673    }
674    fn record_i64(&mut self, field: &tracing::field::Field, value: i64) {
675        self.add_value(field, value.to_string());
676    }
677    fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
678        self.add_value(field, format!("{:?}", value));
679    }
680}
681
682impl tracing::Subscriber for DatadogTracing {
683    fn enabled(&self, metadata: &tracing::Metadata<'_>) -> bool {
684        match self.log_config {
685            Some(ref lc) => log_level_to_trace_level(lc.level) >= *metadata.level(),
686            None => false,
687        }
688    }
689
690    fn new_span(&self, span: &tracing::span::Attributes<'_>) -> tracing::span::Id {
691        let nanos = Utc::now().timestamp_nanos_opt().unwrap_or(0) as u64;
692        let mut new_span_visitor = HashMapVisitor::new();
693        span.record(&mut new_span_visitor);
694        let trace_id = new_span_visitor
695            .fields
696            .remove("trace_id")
697            .and_then(|s| s.parse::<u64>().ok())
698            .unwrap_or(Utc::now().timestamp_nanos_opt().unwrap_or(0) as u64);
699        let span_id = Utc::now().timestamp_nanos_opt().unwrap_or(0) as u64 + 1;
700        let new_span = NewSpanData {
701            id: span_id,
702            trace_id,
703            start: Utc::now(),
704            resource: span.metadata().target().to_string(),
705            name: span.metadata().name().to_string(),
706        };
707        self.send_new_span(nanos, new_span).unwrap_or(());
708        tracing::span::Id::from_u64(span_id)
709    }
710
711    fn record(&self, _span: &tracing::span::Id, _values: &tracing::span::Record<'_>) {}
712
713    fn record_follows_from(&self, _span: &tracing::span::Id, _follows: &tracing::span::Id) {}
714
715    fn event(&self, event: &tracing::Event<'_>) {
716        let thread_id = get_thread_id();
717        let mut new_evt_visitor = HashMapVisitor::new();
718        event.record(&mut new_evt_visitor);
719        self.send_event(
720            thread_id,
721            new_evt_visitor.fields,
722            Utc::now(),
723            event.metadata().level(),
724            event.metadata().module_path().map(|s| s.to_string()),
725        )
726        .unwrap_or(());
727    }
728
729    fn enter(&self, span: &tracing::span::Id) {
730        let nanos = Utc::now().timestamp_nanos_opt().unwrap_or(0) as u64;
731        let thread_id = get_thread_id();
732        self.send_enter_span(nanos, thread_id, span.clone().into_u64())
733            .unwrap_or(());
734        set_current_span_id(Some(span.into_u64()));
735    }
736
737    fn exit(&self, span: &tracing::span::Id) {
738        let nanos = Utc::now().timestamp_nanos_opt().unwrap_or(0) as u64;
739        self.send_exit_span(nanos, span.clone().into_u64())
740            .unwrap_or(());
741        set_current_span_id(None);
742    }
743
744    fn try_close(&self, span: tracing::span::Id) -> bool {
745        let nanos = Utc::now().timestamp_nanos_opt().unwrap_or(0) as u64;
746        self.send_close_span(nanos, span.into_u64()).unwrap_or(());
747        false
748    }
749}
750
751#[cfg(feature = "json")]
752struct KeyValueMap(HashMap<String, String>);
753
754#[cfg(feature = "json")]
755impl<'kvs> kv::VisitSource<'kvs> for KeyValueMap {
756    fn visit_pair(&mut self, key: kv::Key<'kvs>, value: kv::Value<'kvs>) -> Result<(), kv::Error> {
757        self.0.insert(key.to_string(), value.to_string());
758        Ok(())
759    }
760}
761
762#[cfg(feature = "json")]
763fn build_key_value_map<'a>(record: &Record<'a>) -> HashMap<String, String> {
764    let mut visitor = KeyValueMap(HashMap::new());
765    let visit_result = record.key_values().visit(&mut visitor);
766    if let Err(e) = visit_result {
767        println!("Error building key value map: {:?}", e);
768    }
769
770    visitor.0
771}
772
773impl Log for DatadogTracing {
774    fn enabled(&self, metadata: &log::Metadata) -> bool {
775        if let Some(ref lc) = self.log_config {
776            metadata.level() <= lc.level
777        } else {
778            false
779        }
780    }
781
782    fn log(&self, record: &Record) {
783        if let Some(ref lc) = self.log_config {
784            #[cfg(feature = "json")]
785            let key_values = build_key_value_map(record);
786            if record.level() <= lc.level {
787                let thread_id = get_thread_id();
788                let now = chrono::Utc::now();
789                let msg_str = format!("{}", record.args());
790                let log_rec = LogRecord {
791                    thread_id,
792                    level: log_level_to_trace_level(record.level()),
793                    time: now,
794                    module: record.module_path().map(|s| s.to_string()),
795                    msg_str,
796                    #[cfg(feature = "json")]
797                    key_values,
798                };
799                self.send_log(log_rec).unwrap_or(());
800            }
801        }
802    }
803
804    fn flush(&self) {}
805}
806
807#[derive(Debug, Clone)]
808struct DdAgentClient {
809    client_sender: crossbeam_channel::Sender<Vec<Span>>,
810}
811
812impl DdAgentClient {
813    fn new(config: &Config) -> Self {
814        let (client_sender, client_requests) = crossbeam_channel::unbounded();
815
816        for _ in 0..config.num_client_send_threads {
817            let env = config.env.clone();
818            let service = config.service.clone();
819            let host = config.host.clone();
820            let port = config.port.clone();
821            let apm_config = config.apm_config.clone();
822            let cr_channel = client_requests.clone();
823            std::thread::spawn(move || {
824                DdAgentClient::thread_loop(
825                    cr_channel,
826                    env,
827                    format!("http://{}:{}/v0.3/traces", host, port),
828                    service,
829                    apm_config,
830                )
831            });
832        }
833        DdAgentClient { client_sender }
834    }
835
836    fn send(&self, stack: Vec<Span>) {
837        self.client_sender.send(stack).unwrap_or_else(|_| {
838            println!("Tracing send error: Channel closed!");
839        });
840    }
841
842    fn thread_loop(
843        client_requests: crossbeam_channel::Receiver<Vec<Span>>,
844        env: Option<String>,
845        endpoint: String,
846        service: String,
847        apm_config: ApmConfig,
848    ) {
849        // Loop as long as the channel is open
850        while let Ok(stack) = client_requests.recv() {
851            let count = stack.len();
852            let spans: Vec<Vec<RawSpan>> = vec![stack
853                .into_iter()
854                .map(|s| RawSpan::from_span(&s, &service, &env, &apm_config))
855                .collect()];
856            match to_string(&spans) {
857                Err(e) => println!("Couldn't encode payload for datadog: {:?}", e),
858                Ok(payload) => {
859                    let req = attohttpc::post(&endpoint)
860                        .header("Content-Length", payload.len() as u64)
861                        .header("Content-Type", "application/json")
862                        .header("X-Datadog-Trace-Count", count)
863                        .text(&payload);
864
865                    match req.send() {
866                        Ok(resp) if !resp.is_success() => {
867                            println!("error from datadog agent: {:?}", resp)
868                        }
869                        Err(err) => println!("error sending traces to datadog: {:?}", err),
870                        _ => {}
871                    }
872                }
873            }
874        }
875    }
876}
877
878#[cfg(test)]
879mod tests {
880    use super::*;
881    use log::{debug, info};
882    use tracing::level_filters;
883    use tracing::{event, span};
884
885    fn long_call(trace_id: u64) {
886        let span = span!(tracing::Level::INFO, "long_call", trace_id = trace_id);
887        let _e = span.enter();
888        debug!("Waiting on I/O {}", trace_id);
889        sleep_call(trace_id);
890        info!("I/O Finished {}", trace_id);
891    }
892
893    fn sleep_call(trace_id: u64) {
894        let span = span!(tracing::Level::INFO, "sleep_call", trace_id = trace_id);
895        let _e = span.enter();
896        debug!("Long call {}", trace_id);
897        debug!(
898            "Current thread ID/span ID: {}/{:?}",
899            get_thread_id(),
900            get_current_span_id()
901        );
902        std::thread::sleep(std::time::Duration::from_millis(2000));
903    }
904
905    fn traced_func_no_send(trace_id: u64) {
906        let span = span!(
907            tracing::Level::INFO,
908            "traced_func_no_send",
909            trace_id = trace_id
910        );
911        let _e = span.enter();
912        debug!(
913            "Performing some function for id={}/{:?}",
914            trace_id,
915            get_current_span_id()
916        );
917        long_call(trace_id);
918    }
919
920    fn traced_http_func(trace_id: u64) {
921        let span = span!(
922            tracing::Level::INFO,
923            "traced_http_func",
924            trace_id = trace_id
925        );
926        let _e = span.enter();
927        debug!(
928            "Performing some function for id={}/{:?}",
929            trace_id,
930            get_current_span_id()
931        );
932        long_call(trace_id);
933        event!(
934            tracing::Level::INFO,
935            http.url = "http://test.test/",
936            http.status_code = "200",
937            http.method = "GET"
938        );
939        event!(tracing::Level::INFO, send_trace = trace_id);
940    }
941
942    fn traced_error_func(trace_id: u64) {
943        let span = span!(
944            tracing::Level::INFO,
945            "traced_error_func",
946            trace_id = trace_id
947        );
948        let _e = span.enter();
949        debug!(
950            "Performing some function for id={}/{:?}",
951            trace_id,
952            get_current_span_id()
953        );
954        long_call(trace_id);
955        event!(
956            tracing::Level::ERROR,
957            error.etype = "",
958            error.message = "Test error"
959        );
960        event!(
961            tracing::Level::ERROR,
962            http.url = "http://test.test/",
963            http.status_code = "400",
964            http.method = "GET"
965        );
966        event!(
967            tracing::Level::ERROR,
968            custom_tag = "good",
969            custom_tag2 = "test",
970            send_trace = trace_id
971        );
972    }
973
974    fn traced_error_func_single_event(trace_id: u64) {
975        let span = span!(
976            tracing::Level::INFO,
977            "traced_error_func_single_event",
978            trace_id = trace_id
979        );
980        let _e = span.enter();
981
982        debug!(
983            "Performing some function for id={}/{:?}",
984            trace_id,
985            get_current_span_id()
986        );
987        long_call(trace_id);
988        event!(
989            tracing::Level::ERROR,
990            send_trace = trace_id,
991            error.etype = "",
992            error.message = "Test error",
993            http.url = "http://test.test/",
994            http.status_code = "400",
995            http.method = "GET",
996            custom_tag = "good",
997            custom_tag2 = "test"
998        );
999    }
1000
1001    fn trace_config(log_level: log::Level, trace_level: tracing::Level) {
1002        let config = Config {
1003            service: String::from("datadog_apm_test"),
1004            env: Some("staging-01".into()),
1005            logging_config: Some(LoggingConfig {
1006                level: log_level,
1007                mod_filter: vec!["hyper", "mime", "test_log"],
1008                ..LoggingConfig::default()
1009            }),
1010            enable_tracing: true,
1011            ..Default::default()
1012        };
1013        let _client = DatadogTracing::new(config)
1014            .with(
1015                filter::targets::Targets::new()
1016                    .with_target(
1017                        "datadog_apm_sync::client::tests::test_trace",
1018                        level_filters::LevelFilter::OFF,
1019                    )
1020                    .with_default(trace_level),
1021            )
1022            .try_init()
1023            .or_else(|e| {
1024                log::warn!("Error initializing logger: {e}");
1025                Result::<(), ()>::Ok(())
1026            })
1027            .unwrap();
1028    }
1029
1030    #[test]
1031    fn test_exit_child_span() {
1032        trace_config(log::Level::Trace, tracing::Level::TRACE);
1033        let trace_id = 1u64;
1034
1035        let f1: std::thread::JoinHandle<()> = std::thread::spawn(move || {
1036            let span = span!(tracing::Level::INFO, "parent_span", trace_id = trace_id);
1037            let _e = span.enter();
1038            info!("Inside parent_span, should print trace and span ID");
1039            {
1040                let span = span!(tracing::Level::INFO, "child_span", trace_id = trace_id);
1041                let _e = span.enter();
1042                info!("Inside child_span, should print trace and span ID");
1043            }
1044            info!("Back in parent_span, should print trace and span ID");
1045        });
1046        f1.join().unwrap();
1047        event!(tracing::Level::INFO, send_trace = trace_id);
1048        ::std::thread::sleep(::std::time::Duration::from_millis(1000));
1049    }
1050
1051    #[test]
1052    fn test_trace_one_func_stack() {
1053        let trace_id = create_unique_id64();
1054        trace_config(log::Level::Trace, tracing::Level::TRACE);
1055
1056        debug!(
1057            "Outside of span, this should be None: {:?}",
1058            get_current_span_id()
1059        );
1060        debug!(
1061            "Sampling rate is {}",
1062            DatadogTracing::get_global_sampling_rate()
1063        );
1064
1065        let f1 = std::thread::spawn(move || {
1066            traced_func_no_send(trace_id);
1067            event!(tracing::Level::INFO, send_trace = trace_id);
1068        });
1069
1070        debug!(
1071            "Same as before span, after span completes, this should be None: {:?}",
1072            get_current_span_id()
1073        );
1074        f1.join().unwrap();
1075        ::std::thread::sleep(::std::time::Duration::from_millis(1000));
1076    }
1077
1078    #[test]
1079    fn test_parallel_two_threads_two_traces() {
1080        let trace_id1 = create_unique_id64();
1081        let trace_id2 = create_unique_id64();
1082        trace_config(log::Level::Trace, tracing::Level::TRACE);
1083        let f1 = std::thread::spawn(move || {
1084            traced_func_no_send(trace_id1);
1085            event!(tracing::Level::INFO, send_trace = trace_id1);
1086        });
1087        let f2 = std::thread::spawn(move || {
1088            traced_func_no_send(trace_id2);
1089            event!(tracing::Level::INFO, send_trace = trace_id2);
1090        });
1091
1092        f1.join().unwrap();
1093        f2.join().unwrap();
1094        ::std::thread::sleep(::std::time::Duration::from_millis(1000));
1095    }
1096
1097    #[test]
1098    fn test_parallel_two_threads_ten_traces() {
1099        let trace_id1 = create_unique_id64();
1100        let trace_id2 = create_unique_id64() + 1;
1101        let trace_id3 = create_unique_id64() + 2;
1102        let trace_id4 = create_unique_id64() + 3;
1103        let trace_id5 = create_unique_id64() + 4;
1104        let trace_id6 = create_unique_id64() + 5;
1105        let trace_id7 = create_unique_id64() + 6;
1106        let trace_id8 = create_unique_id64() + 7;
1107        let trace_id9 = create_unique_id64() + 8;
1108        let trace_id10 = create_unique_id64() + 9;
1109        trace_config(log::Level::Trace, tracing::Level::TRACE);
1110        let f1 = std::thread::spawn(move || {
1111            traced_func_no_send(trace_id1);
1112            event!(tracing::Level::INFO, send_trace = trace_id1);
1113        });
1114        let f2 = std::thread::spawn(move || {
1115            traced_func_no_send(trace_id2);
1116            event!(tracing::Level::INFO, send_trace = trace_id2);
1117        });
1118        let f3 = std::thread::spawn(move || {
1119            traced_func_no_send(trace_id3);
1120            event!(tracing::Level::INFO, send_trace = trace_id3);
1121        });
1122        let f4 = std::thread::spawn(move || {
1123            traced_func_no_send(trace_id4);
1124            event!(tracing::Level::INFO, send_trace = trace_id4);
1125        });
1126        let f5 = std::thread::spawn(move || {
1127            traced_func_no_send(trace_id5);
1128            event!(tracing::Level::INFO, send_trace = trace_id5);
1129        });
1130        let f6 = std::thread::spawn(move || {
1131            traced_func_no_send(trace_id6);
1132            event!(tracing::Level::INFO, send_trace = trace_id6);
1133        });
1134        let f7 = std::thread::spawn(move || {
1135            traced_func_no_send(trace_id7);
1136            event!(tracing::Level::INFO, send_trace = trace_id7);
1137        });
1138        let f8 = std::thread::spawn(move || {
1139            traced_func_no_send(trace_id8);
1140            event!(tracing::Level::INFO, send_trace = trace_id8);
1141        });
1142        let f9 = std::thread::spawn(move || {
1143            traced_func_no_send(trace_id9);
1144            event!(tracing::Level::INFO, send_trace = trace_id9);
1145        });
1146        let f10 = std::thread::spawn(move || {
1147            traced_func_no_send(trace_id10);
1148            event!(tracing::Level::INFO, send_trace = trace_id10);
1149        });
1150        f1.join().unwrap();
1151        f2.join().unwrap();
1152        f3.join().unwrap();
1153        f4.join().unwrap();
1154        f5.join().unwrap();
1155        f6.join().unwrap();
1156        f7.join().unwrap();
1157        f8.join().unwrap();
1158        f9.join().unwrap();
1159        f10.join().unwrap();
1160
1161        ::std::thread::sleep(::std::time::Duration::from_millis(1000));
1162    }
1163
1164    #[test]
1165    fn test_error_span() {
1166        let trace_id = create_unique_id64();
1167        trace_config(log::Level::Trace, tracing::Level::TRACE);
1168        let f3 = std::thread::spawn(move || {
1169            traced_error_func(trace_id);
1170        });
1171        f3.join().unwrap();
1172        ::std::thread::sleep(::std::time::Duration::from_millis(1000));
1173    }
1174
1175    #[test]
1176    fn test_error_span_as_single_event() {
1177        let trace_id = create_unique_id64();
1178        trace_config(log::Level::Trace, tracing::Level::TRACE);
1179        let f4 = std::thread::spawn(move || {
1180            traced_error_func_single_event(trace_id);
1181        });
1182        f4.join().unwrap();
1183        ::std::thread::sleep(::std::time::Duration::from_millis(1000));
1184    }
1185
1186    #[test]
1187    fn test_two_funcs_in_one_span() {
1188        let trace_id = create_unique_id64();
1189        trace_config(log::Level::Trace, tracing::Level::TRACE);
1190        let f5 = std::thread::spawn(move || {
1191            traced_func_no_send(trace_id);
1192            traced_func_no_send(trace_id);
1193            // Send both funcs under one parent span and one trace
1194            event!(tracing::Level::INFO, send_trace = trace_id);
1195        });
1196        f5.join().unwrap();
1197        ::std::thread::sleep(::std::time::Duration::from_millis(1000));
1198    }
1199
1200    #[test]
1201    fn test_one_thread_two_funcs_serial_two_traces() {
1202        let trace_id1 = create_unique_id64();
1203        let trace_id2 = create_unique_id64();
1204        trace_config(log::Level::Trace, tracing::Level::TRACE);
1205        let f7 = std::thread::spawn(move || {
1206            traced_func_no_send(trace_id1);
1207            event!(tracing::Level::INFO, send_trace = trace_id1);
1208
1209            traced_func_no_send(trace_id2);
1210            event!(tracing::Level::INFO, send_trace = trace_id2);
1211        });
1212        f7.join().unwrap();
1213        ::std::thread::sleep(::std::time::Duration::from_millis(1000));
1214    }
1215
1216    #[test]
1217    fn test_http_span() {
1218        let trace_id = create_unique_id64();
1219        trace_config(log::Level::Trace, tracing::Level::TRACE);
1220        let f3 = std::thread::spawn(move || {
1221            traced_http_func(trace_id);
1222        });
1223        f3.join().unwrap();
1224        ::std::thread::sleep(::std::time::Duration::from_millis(1000));
1225    }
1226
1227    pub mod test_log {
1228        pub fn test_log_fn() {
1229            tracing::event!(
1230                tracing::Level::INFO,
1231                message = "TEST_INFO EVENT in filtered mod - SHOULD ____NOT____ SEE!!"
1232            );
1233            log::info!("TEST_INFO LOG in filtered mod - SHOULD ____NOT____ SEE!!");
1234        }
1235    }
1236
1237    pub mod test_trace {
1238        pub fn test_trace_fn() {
1239            tracing::event!(
1240                tracing::Level::INFO,
1241                message = "TEST_INFO EVENT in filtered trace mod - SHOULD ____NOT____ SEE!!"
1242            );
1243        }
1244    }
1245
1246    #[test]
1247    fn test_log() {
1248        let _trace_id = create_unique_id64();
1249        trace_config(log::Level::Info, tracing::Level::INFO);
1250        log::info!("TEST_INFO - SHOULD SEE!!");
1251        log::debug!("TEST_DEBUG - SHOULD NOT SEE!!");
1252
1253        test_log::test_log_fn();
1254    }
1255
1256    use tracing_subscriber::filter;
1257    use tracing_subscriber::prelude::*;
1258
1259    #[test]
1260    fn test_trace_event_log() {
1261        let _trace_id = create_unique_id64();
1262        trace_config(log::Level::Info, tracing::Level::INFO);
1263
1264        tracing::info!("TEST_INFO EVENT - SHOULD SEE!!");
1265        tracing::debug!("TEST_DEBUG - SHOULD ____NOT____ SEE!!");
1266        tracing::event!(
1267            tracing::Level::INFO,
1268            message = "TEST_INFO EVENT - SHOULD SEE!!"
1269        );
1270        tracing::event!(
1271            tracing::Level::DEBUG,
1272            message = "TEST_DEBUG EVENT - SHOULD ____NOT____ SEE!!"
1273        );
1274
1275        test_log::test_log_fn();
1276        test_trace::test_trace_fn();
1277    }
1278}