libdd_trace_utils/send_data/
mod.rs

1// Copyright 2024-Present Datadog, Inc. https://www.datadoghq.com/
2// SPDX-License-Identifier: Apache-2.0
3
4pub mod send_data_result;
5
6use crate::msgpack_encoder;
7use crate::send_with_retry::compression::{add_headers, compress};
8use crate::send_with_retry::{
9    send_with_retry, CompressionStrategy, RetryStrategy, SendWithRetryResult,
10};
11use crate::trace_utils::TracerHeaderTags;
12use crate::tracer_payload::TracerPayloadCollection;
13use anyhow::{anyhow, Context};
14use futures::stream::FuturesUnordered;
15use futures::StreamExt;
16use http::{header::CONTENT_TYPE, HeaderMap, HeaderValue};
17use libdd_capabilities::{HttpClientCapability, SleepCapability};
18use libdd_common::{
19    header::{
20        APPLICATION_MSGPACK, APPLICATION_PROTOBUF, DATADOG_SEND_REAL_HTTP_STATUS,
21        DATADOG_TRACE_COUNT,
22    },
23    Endpoint,
24};
25use libdd_trace_protobuf::pb::{AgentPayload, TracerPayload};
26use send_data_result::SendDataResult;
27use std::collections::HashMap;
28
29#[derive(Debug)]
30/// `SendData` is a structure that holds the data to be sent to a target endpoint.
31/// It includes the payloads to be sent, the size of the data, the target endpoint,
32/// headers for the request, and a retry strategy for sending the data.
33///
34/// # Example
35///
36/// ```rust
37/// use libdd_trace_protobuf::pb::TracerPayload;
38/// use libdd_trace_utils::send_data::{
39///     SendData,
40/// };
41/// use libdd_common::Endpoint;
42/// use libdd_trace_utils::send_with_retry::{RetryBackoffType, RetryStrategy};
43/// use libdd_trace_utils::trace_utils::TracerHeaderTags;
44/// use libdd_trace_utils::tracer_payload::TracerPayloadCollection;
45///
46/// #[cfg_attr(miri, ignore)]
47/// async fn update_send_results_example() {
48///     let size = 100;
49///     let tracer_payload = TracerPayloadCollection::V07(
50///         vec![TracerPayload::default()]); // Replace with actual payload
51///     let tracer_header_tags = TracerHeaderTags::default(); // Replace with actual header tags
52///     let target = Endpoint::default(); // Replace with actual endpoint
53///
54///     let mut send_data = SendData::new(size, tracer_payload, tracer_header_tags, &target);
55///
56///     // Set a custom retry strategy
57///     let retry_strategy = RetryStrategy::new(3, 10, RetryBackoffType::Exponential, Some(5));
58///
59///     send_data.set_retry_strategy(retry_strategy);
60///
61///     // Send the data (caller picks the capabilities implementation)
62///     use libdd_capabilities::HttpClientCapability;
63///     let capabilities = libdd_capabilities_impl::NativeCapabilities::new_client();
64///     let result = send_data.send(&capabilities).await;
65/// }
66/// ```
67pub struct SendData {
68    pub(crate) tracer_payloads: TracerPayloadCollection,
69    pub(crate) size: usize, // have a rough size estimate to force flushing if it's large
70    target: Endpoint,
71    headers: HeaderMap,
72    retry_strategy: RetryStrategy,
73    compression: CompressionStrategy,
74}
75
76pub struct SendDataBuilder {
77    pub(crate) tracer_payloads: TracerPayloadCollection,
78    pub(crate) size: usize,
79    target: Endpoint,
80    headers: HeaderMap,
81    retry_strategy: RetryStrategy,
82    compression: CompressionStrategy,
83}
84
85impl SendDataBuilder {
86    pub fn new(
87        size: usize,
88        tracer_payload: TracerPayloadCollection,
89        tracer_header_tags: TracerHeaderTags,
90        target: &Endpoint,
91    ) -> SendDataBuilder {
92        let mut headers: HeaderMap = tracer_header_tags.into();
93        headers.insert(DATADOG_SEND_REAL_HTTP_STATUS, HeaderValue::from_static("1"));
94        SendDataBuilder {
95            tracer_payloads: tracer_payload,
96            size,
97            target: target.clone(),
98            headers,
99            retry_strategy: RetryStrategy::default(),
100            compression: CompressionStrategy::None,
101        }
102    }
103
104    pub fn with_compression(mut self, compression: CompressionStrategy) -> SendDataBuilder {
105        self.compression = compression;
106        self
107    }
108
109    pub fn with_api_key(mut self, api_key: &str) -> SendDataBuilder {
110        self.target.api_key = Some(api_key.to_string().into());
111        self
112    }
113
114    pub fn with_retry_strategy(mut self, retry_strategy: RetryStrategy) -> SendDataBuilder {
115        self.retry_strategy = retry_strategy;
116        self
117    }
118
119    pub fn build(self) -> SendData {
120        SendData {
121            tracer_payloads: self.tracer_payloads,
122            size: self.size,
123            target: self.target,
124            headers: self.headers,
125            retry_strategy: self.retry_strategy,
126            compression: self.compression,
127        }
128    }
129}
130
131impl SendData {
132    /// Creates a new instance of `SendData`.
133    ///
134    /// # Arguments
135    ///
136    /// * `size`: Approximate size of the data to be sent in bytes.
137    /// * `tracer_payload`: The payload to be sent.
138    /// * `tracer_header_tags`: The header tags for the tracer.
139    /// * `target`: The endpoint to which the data will be sent.
140    ///
141    /// # Returns
142    ///
143    /// A new `SendData` instance.
144    #[allow(unused_variables)]
145    pub fn new(
146        size: usize,
147        tracer_payload: TracerPayloadCollection,
148        tracer_header_tags: TracerHeaderTags,
149        target: &Endpoint,
150    ) -> SendData {
151        let mut headers: HeaderMap = tracer_header_tags.into();
152        headers.insert(DATADOG_SEND_REAL_HTTP_STATUS, HeaderValue::from_static("1"));
153        SendData {
154            tracer_payloads: tracer_payload,
155            size,
156            target: target.clone(),
157            headers,
158            retry_strategy: RetryStrategy::default(),
159            compression: CompressionStrategy::None,
160        }
161    }
162
163    /// Returns the user defined approximate size of the data to be sent in bytes.
164    ///
165    /// # Returns
166    ///
167    /// The size of the data.
168    pub fn len(&self) -> usize {
169        self.size
170    }
171
172    /// Checks if the user defined approximate size of the data to be sent is zero.
173    ///
174    /// # Returns
175    ///
176    /// `true` if size is 0, `false` otherwise.
177    pub fn is_empty(&self) -> bool {
178        self.size == 0
179    }
180
181    /// Returns the target endpoint.
182    ///
183    /// # Returns
184    ///
185    /// A reference to the target endpoint.
186    pub fn get_target(&self) -> &Endpoint {
187        &self.target
188    }
189
190    /// Returns the payloads to be sent.
191    ///
192    /// # Returns
193    ///
194    /// A reference to the vector of payloads.
195    pub fn get_payloads(&self) -> &TracerPayloadCollection {
196        &self.tracer_payloads
197    }
198
199    /// Overrides the default RetryStrategy with user-defined values.
200    ///
201    /// # Arguments
202    ///
203    /// * `retry_strategy`: The new retry strategy to be used.
204    pub fn set_retry_strategy(&mut self, retry_strategy: RetryStrategy) {
205        self.retry_strategy = retry_strategy;
206    }
207
208    /// Sends the data to the target endpoint.
209    ///
210    /// # Returns
211    ///
212    /// A `SendDataResult` instance containing the result of the operation.
213    pub async fn send<C: HttpClientCapability + SleepCapability>(
214        &self,
215        capabilities: &C,
216    ) -> SendDataResult {
217        self.send_internal(capabilities, None).await
218    }
219
220    async fn send_internal<C: HttpClientCapability + SleepCapability>(
221        &self,
222        capabilities: &C,
223        endpoint: Option<Endpoint>,
224    ) -> SendDataResult {
225        if self.use_protobuf() {
226            self.send_with_protobuf(capabilities, endpoint).await
227        } else {
228            self.send_with_msgpack(capabilities, endpoint).await
229        }
230    }
231
232    async fn send_payload<C: HttpClientCapability + SleepCapability>(
233        &self,
234        capabilities: &C,
235        chunks: u64,
236        payload: Vec<u8>,
237        mut headers: HeaderMap,
238        endpoint: Option<&Endpoint>,
239        compression_strategy: CompressionStrategy,
240    ) -> (SendWithRetryResult, u64, u64) {
241        // Compress here (rather than inside `send_with_retry`) so that the reported
242        // `bytes_sent` metric reflects the number of bytes actually put on the wire.
243        let (payload, compression_strategy) = compress(payload, compression_strategy);
244        add_headers(&mut headers, compression_strategy);
245        #[allow(clippy::unwrap_used)]
246        let payload_len = u64::try_from(payload.len()).unwrap();
247        (
248            send_with_retry(
249                capabilities,
250                endpoint.unwrap_or(&self.target),
251                payload,
252                &headers,
253                &self.retry_strategy,
254                CompressionStrategy::None,
255            )
256            .await,
257            payload_len,
258            chunks,
259        )
260    }
261
262    fn use_protobuf(&self) -> bool {
263        self.target.api_key.is_some()
264    }
265
266    async fn send_with_protobuf<C: HttpClientCapability + SleepCapability>(
267        &self,
268        capabilities: &C,
269        endpoint: Option<Endpoint>,
270    ) -> SendDataResult {
271        let mut result = SendDataResult::default();
272
273        #[allow(clippy::unwrap_used)]
274        let chunks = u64::try_from(self.tracer_payloads.size()).unwrap();
275
276        match &self.tracer_payloads {
277            TracerPayloadCollection::V07(payloads) => {
278                let agent_payload = construct_agent_payload(payloads.to_vec());
279                let serialized_trace_payload = match serialize_proto_payload(&agent_payload)
280                    .context("Failed to serialize trace agent payload, dropping traces")
281                {
282                    Ok(p) => p,
283                    Err(e) => return result.error(e),
284                };
285                let mut request_headers = self.headers.clone();
286
287                request_headers.insert(CONTENT_TYPE, APPLICATION_PROTOBUF);
288
289                let (response, bytes_sent, chunks) = self
290                    .send_payload(
291                        capabilities,
292                        chunks,
293                        serialized_trace_payload,
294                        request_headers,
295                        endpoint.as_ref(),
296                        self.compression,
297                    )
298                    .await;
299
300                result.update(response, bytes_sent, chunks);
301
302                result
303            }
304            _ => result,
305        }
306    }
307
308    async fn send_with_msgpack<C: HttpClientCapability + SleepCapability>(
309        &self,
310        capabilities: &C,
311        endpoint: Option<Endpoint>,
312    ) -> SendDataResult {
313        let mut result = SendDataResult::default();
314        let mut futures = FuturesUnordered::new();
315
316        match &self.tracer_payloads {
317            TracerPayloadCollection::V07(payloads) => {
318                for tracer_payload in payloads {
319                    #[allow(clippy::unwrap_used)]
320                    let chunks = u64::try_from(tracer_payload.chunks.len()).unwrap();
321                    let mut headers = self.headers.clone();
322                    headers.reserve(2);
323                    headers.insert(DATADOG_TRACE_COUNT, chunks.into());
324                    headers.insert(CONTENT_TYPE, APPLICATION_MSGPACK);
325
326                    let payload = match rmp_serde::to_vec_named(tracer_payload) {
327                        Ok(p) => p,
328                        Err(e) => return result.error(anyhow!(e)),
329                    };
330
331                    futures.push(self.send_payload(
332                        capabilities,
333                        chunks,
334                        payload,
335                        headers,
336                        endpoint.as_ref(),
337                        CompressionStrategy::None,
338                    ));
339                }
340            }
341            TracerPayloadCollection::V04(payload) => {
342                #[allow(clippy::unwrap_used)]
343                let chunks = u64::try_from(self.tracer_payloads.size()).unwrap();
344                let mut headers = self.headers.clone();
345                headers.reserve(2);
346                headers.insert(DATADOG_TRACE_COUNT, chunks.into());
347                headers.insert(CONTENT_TYPE, APPLICATION_MSGPACK);
348
349                let payload = msgpack_encoder::v04::to_vec_from_v04(payload);
350
351                futures.push(self.send_payload(
352                    capabilities,
353                    chunks,
354                    payload,
355                    headers,
356                    endpoint.as_ref(),
357                    CompressionStrategy::None,
358                ));
359            }
360            TracerPayloadCollection::V05(payload) => {
361                #[allow(clippy::unwrap_used)]
362                let chunks = u64::try_from(self.tracer_payloads.size()).unwrap();
363                let mut headers = self.headers.clone();
364                headers.reserve(2);
365                headers.insert(DATADOG_TRACE_COUNT, chunks.into());
366                headers.insert(CONTENT_TYPE, APPLICATION_MSGPACK);
367
368                let payload = match rmp_serde::to_vec(payload) {
369                    Ok(p) => p,
370                    Err(e) => return result.error(anyhow!(e)),
371                };
372
373                futures.push(self.send_payload(
374                    capabilities,
375                    chunks,
376                    payload,
377                    headers,
378                    endpoint.as_ref(),
379                    CompressionStrategy::None,
380                ));
381            }
382            TracerPayloadCollection::V1(payload) => {
383                #[allow(clippy::unwrap_used)]
384                let chunks = u64::try_from(self.tracer_payloads.size()).unwrap();
385                let mut headers = self.headers.clone();
386                headers.reserve(2);
387                headers.insert(DATADOG_TRACE_COUNT, chunks.into());
388                headers.insert(CONTENT_TYPE, APPLICATION_MSGPACK);
389
390                let payload = msgpack_encoder::v1::to_vec_from_v1(payload);
391
392                futures.push(self.send_payload(
393                    capabilities,
394                    chunks,
395                    payload,
396                    headers,
397                    endpoint.as_ref(),
398                    CompressionStrategy::None,
399                ));
400            }
401        }
402
403        loop {
404            match futures.next().await {
405                Some((response, payload_len, chunks)) => {
406                    result.update(response, payload_len, chunks);
407                    if result.last_result.is_err() {
408                        return result;
409                    }
410                }
411                None => return result,
412            }
413        }
414    }
415}
416
417fn construct_agent_payload(tracer_payloads: Vec<TracerPayload>) -> AgentPayload {
418    AgentPayload {
419        host_name: "".to_string(),
420        env: "".to_string(),
421        agent_version: "".to_string(),
422        error_tps: 60.0,
423        target_tps: 60.0,
424        tags: HashMap::new(),
425        tracer_payloads,
426        rare_sampler_enabled: false,
427        idx_tracer_payloads: Vec::new(),
428    }
429}
430
431fn serialize_proto_payload<T>(payload: &T) -> anyhow::Result<Vec<u8>>
432where
433    T: prost::Message,
434{
435    let mut buf = Vec::with_capacity(payload.encoded_len());
436    payload.encode(&mut buf)?;
437    Ok(buf)
438}
439
440#[cfg(test)]
441mod tests {
442    use super::*;
443    use crate::send_with_retry::{RetryBackoffType, RetryStrategy};
444    use crate::test_utils::create_test_no_alloc_span;
445    use crate::trace_utils::{construct_trace_chunk, construct_tracer_payload, TracerPayloadTags};
446    use crate::tracer_header_tags::{TracerGenericTags, TracerHeaderTags};
447    use httpmock::prelude::*;
448    use httpmock::MockServer;
449    use libdd_capabilities::HttpClientCapability;
450    use libdd_capabilities_impl::NativeCapabilities;
451    use libdd_common::Endpoint;
452    use libdd_trace_protobuf::pb::Span;
453    use std::collections::HashMap;
454    use std::time::Duration;
455
456    const ONE_SECOND: u64 = 1_000;
457    const HEADER_TAGS: TracerHeaderTags = TracerHeaderTags {
458        lang: "test-lang",
459        lang_version: "2.0",
460        lang_interpreter: "interpreter",
461        lang_vendor: "vendor",
462        tracer_version: "1.0",
463        container_id: "id",
464        generic: TracerGenericTags {
465            client_computed_top_level: false,
466            client_computed_stats: false,
467            dropped_p0_traces: 0,
468            dropped_p0_spans: 0,
469        },
470    };
471
472    fn setup_payload(header_tags: &TracerHeaderTags) -> TracerPayload {
473        let tracer_payload_tags = TracerPayloadTags {
474            env: "TEST".to_string(),
475            app_version: "1.0".to_string(),
476            hostname: "test_bench".to_string(),
477            runtime_id: "id".to_string(),
478        };
479
480        let chunk = construct_trace_chunk(vec![Span {
481            service: "test-service".to_string(),
482            name: "test-service-name".to_string(),
483            resource: "test-service-resource".to_string(),
484            trace_id: 111,
485            span_id: 222,
486            parent_id: 333,
487            start: 1,
488            duration: 5,
489            error: 0,
490            meta: HashMap::new(),
491            metrics: HashMap::new(),
492            meta_struct: HashMap::new(),
493            r#type: "".to_string(),
494            span_links: vec![],
495            span_events: vec![],
496        }]);
497
498        construct_tracer_payload(vec![chunk], header_tags, tracer_payload_tags)
499    }
500
501    fn compute_payload_len(collection: &TracerPayloadCollection) -> usize {
502        match collection {
503            TracerPayloadCollection::V07(payloads) => {
504                let agent_payload = construct_agent_payload(payloads.to_vec());
505                let serialized_trace_payload = serialize_proto_payload(&agent_payload).unwrap();
506                serialized_trace_payload.len()
507            }
508            _ => 0,
509        }
510    }
511
512    fn rmp_compute_payload_len(collection: &TracerPayloadCollection) -> usize {
513        match collection {
514            TracerPayloadCollection::V07(payloads) => {
515                let mut total: usize = 0;
516                for payload in payloads {
517                    total += rmp_serde::to_vec_named(payload).unwrap().len();
518                }
519                total
520            }
521            TracerPayloadCollection::V04(payloads) => {
522                msgpack_encoder::v04::to_encoded_byte_len_from_v04(payloads) as usize
523            }
524            TracerPayloadCollection::V05(payloads) => rmp_serde::to_vec(payloads).unwrap().len(),
525            TracerPayloadCollection::V1(payload) => {
526                msgpack_encoder::v1::to_encoded_byte_len_from_v1(payload) as usize
527            }
528        }
529    }
530
531    #[test]
532    fn send_data_new_api_key() {
533        let header_tags = TracerHeaderTags::default();
534
535        let payload = setup_payload(&header_tags);
536        let data = SendData::new(
537            100,
538            TracerPayloadCollection::V07(vec![payload]),
539            HEADER_TAGS,
540            &Endpoint {
541                api_key: Some(std::borrow::Cow::Borrowed("TEST-KEY")),
542                url: "/foo/bar?baz".parse::<hyper::Uri>().unwrap(),
543                timeout_ms: ONE_SECOND,
544                ..Endpoint::default()
545            },
546        );
547
548        assert_eq!(data.size, 100);
549
550        assert_eq!(data.target.api_key.unwrap(), "TEST-KEY");
551        assert_eq!(data.target.url.path(), "/foo/bar");
552    }
553
554    #[test]
555    fn send_data_new_no_api_key() {
556        let header_tags = TracerHeaderTags::default();
557
558        let payload = setup_payload(&header_tags);
559        let data = SendData::new(
560            100,
561            TracerPayloadCollection::V07(vec![payload]),
562            header_tags.clone(),
563            &Endpoint {
564                api_key: None,
565                url: "/foo/bar?baz".parse::<hyper::Uri>().unwrap(),
566                timeout_ms: ONE_SECOND,
567                ..Endpoint::default()
568            },
569        );
570
571        assert_eq!(data.size, 100);
572
573        assert_eq!(data.target.api_key, None);
574        assert_eq!(data.target.url.path(), "/foo/bar");
575
576        for (key, value) in &HeaderMap::from(header_tags) {
577            assert_eq!(data.headers.get(key), Some(value));
578        }
579    }
580
581    #[cfg_attr(miri, ignore)]
582    #[tokio::test]
583    async fn request_protobuf() {
584        let server = MockServer::start_async().await;
585
586        let mock = server
587            .mock_async(|when, then| {
588                when.method(POST)
589                    .header("Content-type", "application/x-protobuf")
590                    .header("DD-API-KEY", "TEST-KEY")
591                    .path("/");
592                then.status(202).body("");
593            })
594            .await;
595
596        let header_tags = TracerHeaderTags::default();
597
598        let payload = setup_payload(&header_tags);
599        let data = SendData::new(
600            100,
601            TracerPayloadCollection::V07(vec![payload.clone()]),
602            header_tags,
603            &Endpoint {
604                api_key: Some(std::borrow::Cow::Borrowed("TEST-KEY")),
605                url: server.url("/").parse::<hyper::Uri>().unwrap(),
606                timeout_ms: ONE_SECOND,
607                ..Endpoint::default()
608            },
609        );
610
611        let data_payload_len = compute_payload_len(&data.tracer_payloads);
612        let res = data.send(&NativeCapabilities::new_client()).await;
613
614        mock.assert_async().await;
615
616        assert_eq!(
617            res.last_result.unwrap().status(),
618            http::StatusCode::ACCEPTED
619        );
620        assert_eq!(res.errors_timeout, 0);
621        assert_eq!(res.errors_network, 0);
622        assert_eq!(res.errors_status_code, 0);
623        assert_eq!(res.requests_count, 1);
624        assert_eq!(res.chunks_sent, 1);
625        assert_eq!(res.bytes_sent, data_payload_len as u64);
626        assert_eq!(*res.responses_count_per_code.get(&202).unwrap(), 1_u64);
627    }
628
629    #[cfg_attr(miri, ignore)]
630    #[tokio::test]
631    async fn request_protobuf_several_payloads() {
632        let server = MockServer::start_async().await;
633
634        let mock = server
635            .mock_async(|when, then| {
636                when.method(POST)
637                    .header("Content-type", "application/x-protobuf")
638                    .header("DD-API-KEY", "TEST-KEY")
639                    .path("/");
640                then.status(202).body("");
641            })
642            .await;
643
644        let header_tags = TracerHeaderTags::default();
645
646        let payload = setup_payload(&header_tags);
647        let data = SendData::new(
648            100,
649            TracerPayloadCollection::V07(vec![payload.clone(), payload.clone()]),
650            header_tags,
651            &Endpoint {
652                api_key: Some(std::borrow::Cow::Borrowed("TEST-KEY")),
653                url: server.url("/").parse::<hyper::Uri>().unwrap(),
654                timeout_ms: ONE_SECOND,
655                ..Endpoint::default()
656            },
657        );
658
659        let data_payload_len = compute_payload_len(&data.tracer_payloads);
660        let res = data.send(&NativeCapabilities::new_client()).await;
661
662        mock.assert_async().await;
663
664        assert_eq!(
665            res.last_result.unwrap().status(),
666            http::StatusCode::ACCEPTED
667        );
668        assert_eq!(res.errors_timeout, 0);
669        assert_eq!(res.errors_network, 0);
670        assert_eq!(res.errors_status_code, 0);
671        assert_eq!(res.requests_count, 1);
672        assert_eq!(res.chunks_sent, 2);
673        assert_eq!(res.bytes_sent, data_payload_len as u64);
674        assert_eq!(*res.responses_count_per_code.get(&202).unwrap(), 1_u64);
675    }
676
677    #[cfg_attr(miri, ignore)]
678    #[tokio::test]
679    async fn request_msgpack_v07() {
680        let server = MockServer::start_async().await;
681
682        let header_tags = HEADER_TAGS;
683        let mock = server
684            .mock_async(|when, then| {
685                when.method(POST)
686                    .header(DATADOG_TRACE_COUNT.as_str(), "1")
687                    .header("Content-type", "application/msgpack")
688                    .header("datadog-meta-lang", header_tags.lang)
689                    .header(
690                        "datadog-meta-lang-interpreter",
691                        header_tags.lang_interpreter,
692                    )
693                    .header("datadog-meta-lang-version", header_tags.lang_version)
694                    .header(
695                        "datadog-meta-lang-interpreter-vendor",
696                        header_tags.lang_vendor,
697                    )
698                    .header("datadog-meta-tracer-version", header_tags.tracer_version)
699                    .header("datadog-container-id", header_tags.container_id)
700                    .header("Datadog-Send-Real-Http-Status", "1")
701                    .path("/");
702                then.status(200).body("");
703            })
704            .await;
705
706        let header_tags = HEADER_TAGS;
707
708        let payload = setup_payload(&header_tags);
709        let data = SendData::new(
710            100,
711            TracerPayloadCollection::V07(vec![payload.clone()]),
712            header_tags,
713            &Endpoint {
714                api_key: None,
715                url: server.url("/").parse::<hyper::Uri>().unwrap(),
716                timeout_ms: ONE_SECOND,
717                ..Endpoint::default()
718            },
719        );
720
721        let data_payload_len = rmp_compute_payload_len(&data.tracer_payloads);
722        let res = data.send(&NativeCapabilities::new_client()).await;
723
724        mock.assert_async().await;
725
726        assert_eq!(res.last_result.unwrap().status(), http::StatusCode::OK);
727        assert_eq!(res.errors_timeout, 0);
728        assert_eq!(res.errors_network, 0);
729        assert_eq!(res.errors_status_code, 0);
730        assert_eq!(res.requests_count, 1);
731        assert_eq!(res.chunks_sent, 1);
732        assert_eq!(res.bytes_sent, data_payload_len as u64);
733        assert_eq!(*res.responses_count_per_code.get(&200).unwrap(), 1_u64);
734    }
735
736    #[cfg_attr(miri, ignore)]
737    #[tokio::test]
738    async fn request_msgpack_v04() {
739        let server = MockServer::start_async().await;
740
741        let header_tags = HEADER_TAGS;
742        let mock = server
743            .mock_async(|when, then| {
744                when.method(POST)
745                    .header(DATADOG_TRACE_COUNT.as_str(), "1")
746                    .header("Content-type", "application/msgpack")
747                    .header("datadog-meta-lang", header_tags.lang)
748                    .header(
749                        "datadog-meta-lang-interpreter",
750                        header_tags.lang_interpreter,
751                    )
752                    .header("datadog-meta-lang-version", header_tags.lang_version)
753                    .header(
754                        "datadog-meta-lang-interpreter-vendor",
755                        header_tags.lang_vendor,
756                    )
757                    .header("datadog-meta-tracer-version", header_tags.tracer_version)
758                    .header("datadog-container-id", header_tags.container_id)
759                    .path("/");
760                then.status(200).body("");
761            })
762            .await;
763
764        let header_tags = HEADER_TAGS;
765
766        let trace = vec![create_test_no_alloc_span(1234, 12342, 12341, 1, false)];
767        let data = SendData::new(
768            100,
769            TracerPayloadCollection::V04(vec![trace.clone()]),
770            header_tags,
771            &Endpoint {
772                api_key: None,
773                url: server.url("/").parse::<hyper::Uri>().unwrap(),
774                timeout_ms: ONE_SECOND,
775                ..Endpoint::default()
776            },
777        );
778
779        let data_payload_len = rmp_compute_payload_len(&data.tracer_payloads);
780        let res = data.send(&NativeCapabilities::new_client()).await;
781
782        mock.assert_async().await;
783
784        assert_eq!(res.last_result.unwrap().status(), http::StatusCode::OK);
785        assert_eq!(res.errors_timeout, 0);
786        assert_eq!(res.errors_network, 0);
787        assert_eq!(res.errors_status_code, 0);
788        assert_eq!(res.requests_count, 1);
789        assert_eq!(res.chunks_sent, 1);
790        assert_eq!(res.bytes_sent, data_payload_len as u64);
791        assert_eq!(*res.responses_count_per_code.get(&200).unwrap(), 1_u64);
792    }
793
794    #[cfg_attr(miri, ignore)]
795    #[tokio::test]
796    async fn request_msgpack_several_payloads() {
797        let server = MockServer::start_async().await;
798
799        let mock = server
800            .mock_async(|when, then| {
801                when.method(POST)
802                    .header("Content-type", "application/msgpack")
803                    .path("/");
804                then.status(200).body("");
805            })
806            .await;
807
808        let header_tags = TracerHeaderTags::default();
809
810        let payload = setup_payload(&header_tags);
811        let data = SendData::new(
812            100,
813            TracerPayloadCollection::V07(vec![payload.clone(), payload.clone()]),
814            header_tags,
815            &Endpoint {
816                api_key: None,
817                url: server.url("/").parse::<hyper::Uri>().unwrap(),
818                timeout_ms: ONE_SECOND,
819                ..Endpoint::default()
820            },
821        );
822
823        let data_payload_len = rmp_compute_payload_len(&data.tracer_payloads);
824        let res = data.send(&NativeCapabilities::new_client()).await;
825
826        mock.assert_calls_async(2).await;
827
828        assert_eq!(res.last_result.unwrap().status(), http::StatusCode::OK);
829        assert_eq!(res.errors_timeout, 0);
830        assert_eq!(res.errors_network, 0);
831        assert_eq!(res.errors_status_code, 0);
832        assert_eq!(res.requests_count, 2);
833        assert_eq!(res.chunks_sent, 2);
834        assert_eq!(res.bytes_sent, data_payload_len as u64);
835        assert_eq!(*res.responses_count_per_code.get(&200).unwrap(), 2_u64);
836    }
837
838    #[cfg_attr(miri, ignore)]
839    #[tokio::test]
840    async fn request_error_status_code() {
841        let server = MockServer::start_async().await;
842
843        let mock = server
844            .mock_async(|when, then| {
845                when.method(POST)
846                    .header("Content-type", "application/msgpack")
847                    .path("/");
848                then.status(500).body("");
849            })
850            .await;
851
852        let payload = setup_payload(&HEADER_TAGS);
853        let data = SendData::new(
854            100,
855            TracerPayloadCollection::V07(vec![payload]),
856            HEADER_TAGS,
857            &Endpoint {
858                api_key: None,
859                url: server.url("/").parse::<hyper::Uri>().unwrap(),
860                timeout_ms: ONE_SECOND,
861                ..Endpoint::default()
862            },
863        );
864
865        let res = data.send(&NativeCapabilities::new_client()).await;
866
867        mock.assert_calls_async(6).await;
868
869        assert!(res.last_result.is_ok());
870        assert_eq!(
871            res.last_result.unwrap().status(),
872            http::StatusCode::INTERNAL_SERVER_ERROR
873        );
874        assert_eq!(res.errors_timeout, 0);
875        assert_eq!(res.errors_network, 0);
876        assert_eq!(res.errors_status_code, 1);
877        assert_eq!(res.requests_count, 6);
878        assert_eq!(res.chunks_sent, 0);
879        assert_eq!(res.bytes_sent, 0);
880        assert_eq!(*res.responses_count_per_code.get(&500).unwrap(), 1_u64);
881    }
882
883    #[cfg_attr(miri, ignore)]
884    #[tokio::test]
885    async fn request_error_network() {
886        // Server not created in order to return a 'connection refused' error.
887        let payload = setup_payload(&HEADER_TAGS);
888        let data = SendData::new(
889            100,
890            TracerPayloadCollection::V07(vec![payload]),
891            HEADER_TAGS,
892            &Endpoint {
893                api_key: None,
894                url: "http://127.0.0.1:4321/".parse::<hyper::Uri>().unwrap(),
895                timeout_ms: ONE_SECOND,
896                ..Endpoint::default()
897            },
898        );
899
900        let res = data.send(&NativeCapabilities::new_client()).await;
901
902        assert!(res.last_result.is_err());
903        match std::env::consts::OS {
904            "windows" => {
905                // On windows the TCP/IP stack returns a timeout error (at hyper level) rather
906                // than a connection refused error despite not having a listening socket on the
907                // port.
908                assert_eq!(res.errors_timeout, 1);
909                assert_eq!(res.errors_network, 0);
910            }
911            _ => {
912                assert_eq!(res.errors_timeout, 0);
913                assert_eq!(res.errors_network, 1);
914            }
915        }
916        assert_eq!(res.errors_status_code, 0);
917        assert_eq!(res.requests_count, 6);
918        assert_eq!(res.errors_status_code, 0);
919        assert_eq!(res.chunks_sent, 0);
920        assert_eq!(res.bytes_sent, 0);
921        assert_eq!(res.responses_count_per_code.len(), 0);
922    }
923
924    #[cfg_attr(miri, ignore)]
925    #[tokio::test]
926    async fn request_error_timeout_v04() {
927        let server = MockServer::start_async().await;
928
929        let header_tags = HEADER_TAGS;
930        let mock = server
931            .mock_async(|when, then| {
932                when.method(POST)
933                    .header(DATADOG_TRACE_COUNT.as_str(), "2")
934                    .header("Content-type", "application/msgpack")
935                    .header("datadog-meta-lang", header_tags.lang)
936                    .header(
937                        "datadog-meta-lang-interpreter",
938                        header_tags.lang_interpreter,
939                    )
940                    .header("datadog-meta-lang-version", header_tags.lang_version)
941                    .header(
942                        "datadog-meta-lang-interpreter-vendor",
943                        header_tags.lang_vendor,
944                    )
945                    .header("datadog-meta-tracer-version", header_tags.tracer_version)
946                    .header("datadog-container-id", header_tags.container_id)
947                    .path("/");
948                then.status(200)
949                    .body("")
950                    .delay(Duration::from_millis(1_000));
951            })
952            .await;
953
954        let header_tags = HEADER_TAGS;
955
956        let trace = vec![create_test_no_alloc_span(1234, 12342, 12341, 1, false)];
957        let mut data = SendData::new(
958            100,
959            TracerPayloadCollection::V04(vec![trace.clone(), trace.clone()]),
960            header_tags,
961            &Endpoint {
962                api_key: None,
963                url: server.url("/").parse::<hyper::Uri>().unwrap(),
964                timeout_ms: 500,
965                ..Endpoint::default()
966            },
967        );
968
969        // Setting constant time in order to reduce wall time in CI tests in case there is too much
970        // contention.
971        data.set_retry_strategy(RetryStrategy::new(5, 100, RetryBackoffType::Constant, None));
972
973        let res = data.send(&NativeCapabilities::new_client()).await;
974
975        mock.assert_calls_async(6).await;
976
977        assert_eq!(res.errors_timeout, 1);
978        assert_eq!(res.errors_network, 0);
979        assert_eq!(res.errors_status_code, 0);
980        assert_eq!(res.requests_count, 6);
981        assert_eq!(res.chunks_sent, 0);
982        assert_eq!(res.bytes_sent, 0);
983        assert_eq!(res.responses_count_per_code.len(), 0);
984    }
985
986    #[cfg_attr(miri, ignore)]
987    #[tokio::test]
988    async fn request_error_timeout_v07() {
989        let server = MockServer::start_async().await;
990
991        let mock = server
992            .mock_async(|when, then| {
993                when.method(POST)
994                    .header("Content-type", "application/msgpack")
995                    .path("/");
996                then.status(200)
997                    .body("")
998                    .delay(Duration::from_millis(1_000));
999            })
1000            .await;
1001
1002        let header_tags = TracerHeaderTags::default();
1003
1004        let payload = setup_payload(&header_tags);
1005        let mut data = SendData::new(
1006            100,
1007            TracerPayloadCollection::V07(vec![payload.clone(), payload.clone()]),
1008            header_tags,
1009            &Endpoint {
1010                api_key: None,
1011                url: server.url("/").parse::<hyper::Uri>().unwrap(),
1012                timeout_ms: 500,
1013                ..Endpoint::default()
1014            },
1015        );
1016
1017        // Setting constant time in order to reduce wall time in CI tests in case there is too much
1018        // contention.
1019        data.set_retry_strategy(RetryStrategy::new(5, 100, RetryBackoffType::Constant, None));
1020
1021        let res = data.send(&NativeCapabilities::new_client()).await;
1022
1023        mock.assert_calls_async(12).await;
1024
1025        assert_eq!(res.errors_timeout, 1);
1026        assert_eq!(res.errors_network, 0);
1027        assert_eq!(res.errors_status_code, 0);
1028        assert_eq!(res.requests_count, 6);
1029        assert_eq!(res.chunks_sent, 0);
1030        assert_eq!(res.bytes_sent, 0);
1031        assert_eq!(res.responses_count_per_code.len(), 0);
1032    }
1033
1034    #[test]
1035    fn test_builder() {
1036        let header_tags = HEADER_TAGS;
1037        let payload = setup_payload(&header_tags);
1038        let retry_strategy = RetryStrategy::new(5, 100, RetryBackoffType::Constant, None);
1039
1040        let send_data = SendDataBuilder::new(
1041            100,
1042            TracerPayloadCollection::V07(vec![payload]),
1043            header_tags,
1044            &Endpoint::default(),
1045        )
1046        // Test with_api_key()
1047        .with_api_key("TEST-KEY")
1048        // Test with_retry_strategy()
1049        .with_retry_strategy(retry_strategy.clone())
1050        .build();
1051
1052        assert_eq!(
1053            send_data.target.api_key,
1054            Some(std::borrow::Cow::Borrowed("TEST-KEY"))
1055        );
1056        assert_eq!(send_data.retry_strategy, retry_strategy);
1057    }
1058}
libdd_trace_utils/send_data/mod.rs

libdd_trace_utils/send_data/
mod.rs
