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#[derive(Clone, Debug)]
22pub struct Config {
23 pub service: String,
25 pub env: Option<String>,
27 pub host: String,
29 pub port: String,
31 pub logging_config: Option<LoggingConfig>,
33 pub apm_config: ApmConfig,
35 pub enable_tracing: bool,
37 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 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 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 fn enter_span(&mut self, span_id: SpanId) {
203 self.entered_spans.push_back(span_id);
204 }
205
206 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 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 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 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 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 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 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 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 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 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 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 if let Some(send_trace_id) = event.send_trace {
441 let send_vec = storage.drain_completed(send_trace_id, event.time);
442 storage.remove_current_trace(send_trace_id);
445 if !send_vec.is_empty() {
446 client.send(send_vec);
447 }
448 }
449
450 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 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
635pub 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 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 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 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}