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::{warn, 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 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 !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: HashMap::new(),
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            tracing::subscriber::set_global_default(tracer.clone()).unwrap_or_else(|_| {
542                warn!(
543                    "Global subscriber has already been set!  \
544                           This should only be set once in the executable."
545                )
546            });
547        }
548        tracer
549    }
550
551    pub fn get_global_sampling_rate() -> f64 {
552        SAMPLING_RATE.load(Ordering::Acquire)
553    }
554
555    fn send_log(&self, record: LogRecord) -> Result<(), ()> {
556        self.buffer_sender
557            .send(TraceCommand::Log(record))
558            .map(|_| ())
559            .map_err(|_| ())
560    }
561
562    fn send_new_span(&self, nanos: u64, span: NewSpanData) -> Result<(), ()> {
563        self.buffer_sender
564            .send(TraceCommand::NewSpan(nanos, span))
565            .map(|_| ())
566            .map_err(|_| ())
567    }
568
569    fn send_enter_span(&self, nanos: u64, thread_id: ThreadId, id: SpanId) -> Result<(), ()> {
570        self.buffer_sender
571            .send(TraceCommand::Enter(nanos, thread_id, id))
572            .map(|_| ())
573            .map_err(|_| ())
574    }
575
576    fn send_exit_span(&self, nanos: u64, id: SpanId) -> Result<(), ()> {
577        self.buffer_sender
578            .send(TraceCommand::Exit(nanos, id))
579            .map(|_| ())
580            .map_err(|_| ())
581    }
582
583    fn send_close_span(&self, nanos: u64, span_id: SpanId) -> Result<(), ()> {
584        self.buffer_sender
585            .send(TraceCommand::CloseSpan(nanos, span_id))
586            .map(|_| ())
587            .map_err(|_| ())
588    }
589
590    fn send_event(
591        &self,
592        thread_id: ThreadId,
593        event: HashMap<String, String>,
594        time: DateTime<Utc>,
595        level: &tracing::Level,
596        module: Option<String>,
597    ) -> Result<(), ()> {
598        self.buffer_sender
599            .send(TraceCommand::Event(EventRecord::new(
600                thread_id, event, time, *level, module,
601            )))
602            .map(|_| ())
603            .map_err(|_| ())
604    }
605}
606
607fn log_level_to_trace_level(level: log::Level) -> tracing::Level {
608    use log::Level::*;
609    match level {
610        Error => tracing::Level::INFO,
611        Warn => tracing::Level::INFO,
612        Info => tracing::Level::INFO,
613        Debug => tracing::Level::DEBUG,
614        Trace => tracing::Level::TRACE,
615    }
616}
617
618static UNIQUEID_COUNTER: AtomicU16 = AtomicU16::new(0);
619static THREAD_COUNTER: AtomicU32 = AtomicU32::new(0);
620
621static SAMPLING_RATE: AtomicF64 = AtomicF64::new(0.0);
622
623thread_local! {
624    static THREAD_ID: ThreadId = THREAD_COUNTER.fetch_add(1, Ordering::Relaxed);
625    static CURRENT_SPAN_ID: Cell<Option<SpanId>> = const { Cell::new(None) }
626}
627
628pub fn get_thread_id() -> ThreadId {
629    THREAD_ID.with(|id| *id)
630}
631
632pub fn get_current_span_id() -> Option<SpanId> {
633    CURRENT_SPAN_ID.with(|id| id.get())
634}
635
636pub fn set_current_span_id(new_id: Option<SpanId>) {
637    CURRENT_SPAN_ID.with(|id| {
638        id.set(new_id);
639    })
640}
641
642// Format
643// |                       6 bytes                       |      2 bytes    |
644// +--------+--------+--------+--------+--------+--------+--------+--------+
645// |     number of milliseconds since epoch (1970)       | static counter  |
646// +--------+--------+--------+--------+--------+--------+--------+--------+
647// 0        8        16       24       32       40       48       56       64
648//
649// This will hold up to the year 10,000 before it cycles.
650pub fn create_unique_id64() -> u64 {
651    let millis_since_epoch = (Utc::now().timestamp_millis() << 16) as u64;
652    millis_since_epoch + UNIQUEID_COUNTER.fetch_add(1, Ordering::Relaxed) as u64
653}
654
655pub struct HashMapVisitor {
656    fields: HashMap<String, String>,
657}
658
659impl HashMapVisitor {
660    fn new() -> Self {
661        // Event/Span vectors should never have more than ten fields.
662        HashMapVisitor {
663            fields: HashMap::new(),
664        }
665    }
666    fn add_value(&mut self, field: &tracing::field::Field, value: String) {
667        self.fields.insert(field.name().to_string(), value);
668    }
669}
670
671impl tracing::field::Visit for HashMapVisitor {
672    fn record_str(&mut self, field: &tracing::field::Field, value: &str) {
673        self.add_value(field, value.to_string());
674    }
675    fn record_bool(&mut self, field: &tracing::field::Field, value: bool) {
676        self.add_value(field, value.to_string());
677    }
678    fn record_u64(&mut self, field: &tracing::field::Field, value: u64) {
679        self.add_value(field, value.to_string());
680    }
681    fn record_i64(&mut self, field: &tracing::field::Field, value: i64) {
682        self.add_value(field, value.to_string());
683    }
684    fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
685        self.add_value(field, format!("{:?}", value));
686    }
687}
688
689impl tracing::Subscriber for DatadogTracing {
690    fn enabled(&self, metadata: &tracing::Metadata<'_>) -> bool {
691        match self.log_config {
692            Some(ref lc) => log_level_to_trace_level(lc.level) >= *metadata.level(),
693            None => false,
694        }
695    }
696
697    fn new_span(&self, span: &tracing::span::Attributes<'_>) -> tracing::span::Id {
698        let nanos = Utc::now().timestamp_nanos_opt().unwrap_or(0) as u64;
699        let mut new_span_visitor = HashMapVisitor::new();
700        span.record(&mut new_span_visitor);
701        let trace_id = new_span_visitor
702            .fields
703            .remove("trace_id")
704            .and_then(|s| s.parse::<u64>().ok())
705            .unwrap_or(Utc::now().timestamp_nanos_opt().unwrap_or(0) as u64);
706        let span_id = Utc::now().timestamp_nanos_opt().unwrap_or(0) as u64 + 1;
707        let new_span = NewSpanData {
708            id: span_id,
709            trace_id,
710            start: Utc::now(),
711            resource: span.metadata().target().to_string(),
712            name: span.metadata().name().to_string(),
713        };
714        self.send_new_span(nanos, new_span).unwrap_or(());
715        tracing::span::Id::from_u64(span_id)
716    }
717
718    fn record(&self, _span: &tracing::span::Id, _values: &tracing::span::Record<'_>) {}
719
720    fn record_follows_from(&self, _span: &tracing::span::Id, _follows: &tracing::span::Id) {}
721
722    fn event(&self, event: &tracing::Event<'_>) {
723        let thread_id = get_thread_id();
724        let mut new_evt_visitor = HashMapVisitor::new();
725        event.record(&mut new_evt_visitor);
726        self.send_event(
727            thread_id,
728            new_evt_visitor.fields,
729            Utc::now(),
730            event.metadata().level(),
731            event.metadata().module_path().map(|s| s.to_string()),
732        )
733        .unwrap_or(());
734    }
735
736    fn enter(&self, span: &tracing::span::Id) {
737        let nanos = Utc::now().timestamp_nanos_opt().unwrap_or(0) as u64;
738        let thread_id = get_thread_id();
739        self.send_enter_span(nanos, thread_id, span.clone().into_u64())
740            .unwrap_or(());
741        set_current_span_id(Some(span.into_u64()));
742    }
743
744    fn exit(&self, span: &tracing::span::Id) {
745        let nanos = Utc::now().timestamp_nanos_opt().unwrap_or(0) as u64;
746        self.send_exit_span(nanos, span.clone().into_u64())
747            .unwrap_or(());
748        set_current_span_id(None);
749    }
750
751    fn try_close(&self, span: tracing::span::Id) -> bool {
752        let nanos = Utc::now().timestamp_nanos_opt().unwrap_or(0) as u64;
753        self.send_close_span(nanos, span.into_u64()).unwrap_or(());
754        false
755    }
756}
757
758#[cfg(feature = "json")]
759struct KeyValueMap(HashMap<String, String>);
760
761#[cfg(feature = "json")]
762impl<'kvs> kv::VisitSource<'kvs> for KeyValueMap {
763    fn visit_pair(&mut self, key: kv::Key<'kvs>, value: kv::Value<'kvs>) -> Result<(), kv::Error> {
764        self.0.insert(key.to_string(), value.to_string());
765        Ok(())
766    }
767}
768
769#[cfg(feature = "json")]
770fn build_key_value_map<'a>(record: &Record<'a>) -> HashMap<String, String> {
771    let mut visitor = KeyValueMap(HashMap::new());
772    let visit_result = record.key_values().visit(&mut visitor);
773    if let Err(e) = visit_result {
774        println!("Error building key value map: {:?}", e);
775    }
776
777    visitor.0
778}
779
780impl Log for DatadogTracing {
781    fn enabled(&self, metadata: &log::Metadata) -> bool {
782        if let Some(ref lc) = self.log_config {
783            metadata.level() <= lc.level
784        } else {
785            false
786        }
787    }
788
789    fn log(&self, record: &Record) {
790        if let Some(ref lc) = self.log_config {
791            #[cfg(feature = "json")]
792            let key_values = build_key_value_map(record);
793            if record.level() <= lc.level {
794                let thread_id = get_thread_id();
795                let now = chrono::Utc::now();
796                let msg_str = format!("{}", record.args());
797                let log_rec = LogRecord {
798                    thread_id,
799                    level: log_level_to_trace_level(record.level()),
800                    time: now,
801                    module: record.module_path().map(|s| s.to_string()),
802                    msg_str,
803                    #[cfg(feature = "json")]
804                    key_values,
805                };
806                self.send_log(log_rec).unwrap_or(());
807            }
808        }
809    }
810
811    fn flush(&self) {}
812}
813
814#[derive(Debug, Clone)]
815struct DdAgentClient {
816    client_sender: crossbeam_channel::Sender<Vec<Span>>,
817}
818
819impl DdAgentClient {
820    fn new(config: &Config) -> Self {
821        let (client_sender, client_requests) = crossbeam_channel::unbounded();
822
823        for _ in 0..config.num_client_send_threads {
824            let env = config.env.clone();
825            let service = config.service.clone();
826            let host = config.host.clone();
827            let port = config.port.clone();
828            let apm_config = config.apm_config.clone();
829            let cr_channel = client_requests.clone();
830            std::thread::spawn(move || {
831                DdAgentClient::thread_loop(
832                    cr_channel,
833                    env,
834                    format!("http://{}:{}/v0.3/traces", host, port),
835                    service,
836                    apm_config,
837                )
838            });
839        }
840        DdAgentClient { client_sender }
841    }
842
843    fn send(&self, stack: Vec<Span>) {
844        self.client_sender.send(stack).unwrap_or_else(|_| {
845            println!("Tracing send error: Channel closed!");
846        });
847    }
848
849    fn thread_loop(
850        client_requests: crossbeam_channel::Receiver<Vec<Span>>,
851        env: Option<String>,
852        endpoint: String,
853        service: String,
854        apm_config: ApmConfig,
855    ) {
856        // Loop as long as the channel is open
857        while let Ok(stack) = client_requests.recv() {
858            let count = stack.len();
859            let spans: Vec<Vec<RawSpan>> = vec![stack
860                .into_iter()
861                .map(|s| RawSpan::from_span(&s, &service, &env, &apm_config))
862                .collect()];
863            match to_string(&spans) {
864                Err(e) => println!("Couldn't encode payload for datadog: {:?}", e),
865                Ok(payload) => {
866                    let req = attohttpc::post(&endpoint)
867                        .header("Content-Length", payload.len() as u64)
868                        .header("Content-Type", "application/json")
869                        .header("X-Datadog-Trace-Count", count)
870                        .text(&payload);
871
872                    match req.send() {
873                        Ok(resp) if !resp.is_success() => {
874                            println!("error from datadog agent: {:?}", resp)
875                        }
876                        Err(err) => println!("error sending traces to datadog: {:?}", err),
877                        _ => {}
878                    }
879                }
880            }
881        }
882    }
883}
884
885#[cfg(test)]
886mod tests {
887    use super::*;
888    use log::{debug, info, Level};
889    use tracing::{event, span};
890
891    fn long_call(trace_id: u64) {
892        let span = span!(tracing::Level::INFO, "long_call", trace_id = trace_id);
893        let _e = span.enter();
894        debug!("Waiting on I/O {}", trace_id);
895        sleep_call(trace_id);
896        info!("I/O Finished {}", trace_id);
897    }
898
899    fn sleep_call(trace_id: u64) {
900        let span = span!(tracing::Level::INFO, "sleep_call", trace_id = trace_id);
901        let _e = span.enter();
902        debug!("Long call {}", trace_id);
903        debug!(
904            "Current thread ID/span ID: {}/{:?}",
905            get_thread_id(),
906            get_current_span_id()
907        );
908        std::thread::sleep(std::time::Duration::from_millis(2000));
909    }
910
911    fn traced_func_no_send(trace_id: u64) {
912        let span = span!(
913            tracing::Level::INFO,
914            "traced_func_no_send",
915            trace_id = trace_id
916        );
917        let _e = span.enter();
918        debug!(
919            "Performing some function for id={}/{:?}",
920            trace_id,
921            get_current_span_id()
922        );
923        long_call(trace_id);
924    }
925
926    fn traced_http_func(trace_id: u64) {
927        let span = span!(
928            tracing::Level::INFO,
929            "traced_http_func",
930            trace_id = trace_id
931        );
932        let _e = span.enter();
933        debug!(
934            "Performing some function for id={}/{:?}",
935            trace_id,
936            get_current_span_id()
937        );
938        long_call(trace_id);
939        event!(
940            tracing::Level::INFO,
941            http.url = "http://test.test/",
942            http.status_code = "200",
943            http.method = "GET"
944        );
945        event!(tracing::Level::INFO, send_trace = trace_id);
946    }
947
948    fn traced_error_func(trace_id: u64) {
949        let span = span!(
950            tracing::Level::INFO,
951            "traced_error_func",
952            trace_id = trace_id
953        );
954        let _e = span.enter();
955        debug!(
956            "Performing some function for id={}/{:?}",
957            trace_id,
958            get_current_span_id()
959        );
960        long_call(trace_id);
961        event!(
962            tracing::Level::ERROR,
963            error.etype = "",
964            error.message = "Test error"
965        );
966        event!(
967            tracing::Level::ERROR,
968            http.url = "http://test.test/",
969            http.status_code = "400",
970            http.method = "GET"
971        );
972        event!(
973            tracing::Level::ERROR,
974            custom_tag = "good",
975            custom_tag2 = "test",
976            send_trace = trace_id
977        );
978    }
979
980    fn traced_error_func_single_event(trace_id: u64) {
981        let span = span!(
982            tracing::Level::INFO,
983            "traced_error_func_single_event",
984            trace_id = trace_id
985        );
986        let _e = span.enter();
987
988        debug!(
989            "Performing some function for id={}/{:?}",
990            trace_id,
991            get_current_span_id()
992        );
993        long_call(trace_id);
994        event!(
995            tracing::Level::ERROR,
996            send_trace = trace_id,
997            error.etype = "",
998            error.message = "Test error",
999            http.url = "http://test.test/",
1000            http.status_code = "400",
1001            http.method = "GET",
1002            custom_tag = "good",
1003            custom_tag2 = "test"
1004        );
1005    }
1006
1007    fn trace_config() {
1008        let config = Config {
1009            service: String::from("datadog_apm_test"),
1010            env: Some("staging-01".into()),
1011            logging_config: Some(LoggingConfig {
1012                level: Level::Trace,
1013                mod_filter: vec!["hyper", "mime"],
1014                ..LoggingConfig::default()
1015            }),
1016            enable_tracing: true,
1017            ..Default::default()
1018        };
1019        let _client = DatadogTracing::new(config);
1020    }
1021
1022    #[test]
1023    fn test_exit_child_span() {
1024        trace_config();
1025        let trace_id = 1u64;
1026
1027        let f1 = std::thread::spawn(move || {
1028            let span = span!(tracing::Level::INFO, "parent_span", trace_id = trace_id);
1029            let _e = span.enter();
1030            info!("Inside parent_span, should print trace and span ID");
1031            {
1032                let span = span!(tracing::Level::INFO, "child_span", trace_id = trace_id);
1033                let _e = span.enter();
1034                info!("Inside child_span, should print trace and span ID");
1035            }
1036            info!("Back in parent_span, should print trace and span ID");
1037        });
1038        f1.join().unwrap();
1039        event!(tracing::Level::INFO, send_trace = trace_id);
1040        ::std::thread::sleep(::std::time::Duration::from_millis(1000));
1041    }
1042
1043    #[test]
1044    fn test_trace_one_func_stack() {
1045        let trace_id = create_unique_id64();
1046        trace_config();
1047
1048        debug!(
1049            "Outside of span, this should be None: {:?}",
1050            get_current_span_id()
1051        );
1052        debug!(
1053            "Sampling rate is {}",
1054            DatadogTracing::get_global_sampling_rate()
1055        );
1056
1057        let f1 = std::thread::spawn(move || {
1058            traced_func_no_send(trace_id);
1059            event!(tracing::Level::INFO, send_trace = trace_id);
1060        });
1061
1062        debug!(
1063            "Same as before span, after span completes, this should be None: {:?}",
1064            get_current_span_id()
1065        );
1066        f1.join().unwrap();
1067        ::std::thread::sleep(::std::time::Duration::from_millis(1000));
1068    }
1069
1070    #[test]
1071    fn test_parallel_two_threads_two_traces() {
1072        let trace_id1 = create_unique_id64();
1073        let trace_id2 = create_unique_id64();
1074        trace_config();
1075        let f1 = std::thread::spawn(move || {
1076            traced_func_no_send(trace_id1);
1077            event!(tracing::Level::INFO, send_trace = trace_id1);
1078        });
1079        let f2 = std::thread::spawn(move || {
1080            traced_func_no_send(trace_id2);
1081            event!(tracing::Level::INFO, send_trace = trace_id2);
1082        });
1083
1084        f1.join().unwrap();
1085        f2.join().unwrap();
1086        ::std::thread::sleep(::std::time::Duration::from_millis(1000));
1087    }
1088
1089    #[test]
1090    fn test_parallel_two_threads_ten_traces() {
1091        let trace_id1 = create_unique_id64();
1092        let trace_id2 = create_unique_id64() + 1;
1093        let trace_id3 = create_unique_id64() + 2;
1094        let trace_id4 = create_unique_id64() + 3;
1095        let trace_id5 = create_unique_id64() + 4;
1096        let trace_id6 = create_unique_id64() + 5;
1097        let trace_id7 = create_unique_id64() + 6;
1098        let trace_id8 = create_unique_id64() + 7;
1099        let trace_id9 = create_unique_id64() + 8;
1100        let trace_id10 = create_unique_id64() + 9;
1101        trace_config();
1102        let f1 = std::thread::spawn(move || {
1103            traced_func_no_send(trace_id1);
1104            event!(tracing::Level::INFO, send_trace = trace_id1);
1105        });
1106        let f2 = std::thread::spawn(move || {
1107            traced_func_no_send(trace_id2);
1108            event!(tracing::Level::INFO, send_trace = trace_id2);
1109        });
1110        let f3 = std::thread::spawn(move || {
1111            traced_func_no_send(trace_id3);
1112            event!(tracing::Level::INFO, send_trace = trace_id3);
1113        });
1114        let f4 = std::thread::spawn(move || {
1115            traced_func_no_send(trace_id4);
1116            event!(tracing::Level::INFO, send_trace = trace_id4);
1117        });
1118        let f5 = std::thread::spawn(move || {
1119            traced_func_no_send(trace_id5);
1120            event!(tracing::Level::INFO, send_trace = trace_id5);
1121        });
1122        let f6 = std::thread::spawn(move || {
1123            traced_func_no_send(trace_id6);
1124            event!(tracing::Level::INFO, send_trace = trace_id6);
1125        });
1126        let f7 = std::thread::spawn(move || {
1127            traced_func_no_send(trace_id7);
1128            event!(tracing::Level::INFO, send_trace = trace_id7);
1129        });
1130        let f8 = std::thread::spawn(move || {
1131            traced_func_no_send(trace_id8);
1132            event!(tracing::Level::INFO, send_trace = trace_id8);
1133        });
1134        let f9 = std::thread::spawn(move || {
1135            traced_func_no_send(trace_id9);
1136            event!(tracing::Level::INFO, send_trace = trace_id9);
1137        });
1138        let f10 = std::thread::spawn(move || {
1139            traced_func_no_send(trace_id10);
1140            event!(tracing::Level::INFO, send_trace = trace_id10);
1141        });
1142        f1.join().unwrap();
1143        f2.join().unwrap();
1144        f3.join().unwrap();
1145        f4.join().unwrap();
1146        f5.join().unwrap();
1147        f6.join().unwrap();
1148        f7.join().unwrap();
1149        f8.join().unwrap();
1150        f9.join().unwrap();
1151        f10.join().unwrap();
1152
1153        ::std::thread::sleep(::std::time::Duration::from_millis(1000));
1154    }
1155
1156    #[test]
1157    fn test_error_span() {
1158        let trace_id = create_unique_id64();
1159        trace_config();
1160        let f3 = std::thread::spawn(move || {
1161            traced_error_func(trace_id);
1162        });
1163        f3.join().unwrap();
1164        ::std::thread::sleep(::std::time::Duration::from_millis(1000));
1165    }
1166
1167    #[test]
1168    fn test_error_span_as_single_event() {
1169        let trace_id = create_unique_id64();
1170        trace_config();
1171        let f4 = std::thread::spawn(move || {
1172            traced_error_func_single_event(trace_id);
1173        });
1174        f4.join().unwrap();
1175        ::std::thread::sleep(::std::time::Duration::from_millis(1000));
1176    }
1177
1178    #[test]
1179    fn test_two_funcs_in_one_span() {
1180        let trace_id = create_unique_id64();
1181        trace_config();
1182        let f5 = std::thread::spawn(move || {
1183            traced_func_no_send(trace_id);
1184            traced_func_no_send(trace_id);
1185            // Send both funcs under one parent span and one trace
1186            event!(tracing::Level::INFO, send_trace = trace_id);
1187        });
1188        f5.join().unwrap();
1189        ::std::thread::sleep(::std::time::Duration::from_millis(1000));
1190    }
1191
1192    #[test]
1193    fn test_one_thread_two_funcs_serial_two_traces() {
1194        let trace_id1 = create_unique_id64();
1195        let trace_id2 = create_unique_id64();
1196        trace_config();
1197        let f7 = std::thread::spawn(move || {
1198            traced_func_no_send(trace_id1);
1199            event!(tracing::Level::INFO, send_trace = trace_id1);
1200
1201            traced_func_no_send(trace_id2);
1202            event!(tracing::Level::INFO, send_trace = trace_id2);
1203        });
1204        f7.join().unwrap();
1205        ::std::thread::sleep(::std::time::Duration::from_millis(1000));
1206    }
1207
1208    #[test]
1209    fn test_http_span() {
1210        let trace_id = create_unique_id64();
1211        trace_config();
1212        let f3 = std::thread::spawn(move || {
1213            traced_http_func(trace_id);
1214        });
1215        f3.join().unwrap();
1216        ::std::thread::sleep(::std::time::Duration::from_millis(1000));
1217    }
1218}