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#[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 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 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: 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 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
642pub 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 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 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 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}