libdd_trace_utils/
stats_utils.rs

1// Copyright 2023-Present Datadog, Inc. https://www.datadoghq.com/
2// SPDX-License-Identifier: Apache-2.0
3
4#[cfg(feature = "mini_agent")]
5pub use mini_agent::*;
6
7#[cfg(feature = "mini_agent")]
8mod mini_agent {
9    use bytes::{Buf, Bytes};
10    use http_body_util::BodyExt;
11    use libdd_capabilities::HttpClientCapability;
12    use libdd_common::http_common;
13    use libdd_common::Endpoint;
14    use libdd_trace_protobuf::pb;
15    use std::io::Write;
16    use tracing::debug;
17
18    pub async fn get_stats_from_request_body(
19        body: http_common::Body,
20    ) -> anyhow::Result<pb::ClientStatsPayload> {
21        let buffer = BodyExt::collect(body).await?.aggregate();
22
23        let client_stats_payload: pb::ClientStatsPayload =
24            match rmp_serde::from_read(buffer.reader()) {
25                Ok(res) => res,
26                Err(err) => {
27                    anyhow::bail!("Error deserializing stats from request body: {err}")
28                }
29            };
30
31        if client_stats_payload.stats.is_empty() {
32            debug!("Empty trace stats payload received, but this is okay");
33        }
34        Ok(client_stats_payload)
35    }
36
37    pub fn construct_stats_payload(stats: Vec<pb::ClientStatsPayload>) -> pb::StatsPayload {
38        // set hostname on stats from tracer to empty string for serverless
39        let stats = stats
40            .into_iter()
41            .map(|mut stat| {
42                stat.hostname = "".to_string();
43                stat
44            })
45            .collect();
46        pb::StatsPayload {
47            agent_hostname: "".to_string(),
48            agent_env: "".to_string(),
49            stats,
50            agent_version: "".to_string(),
51            client_computed: true,
52            split_payload: false,
53        }
54    }
55
56    pub fn serialize_stats_payload(payload: pb::StatsPayload) -> anyhow::Result<Vec<u8>> {
57        let msgpack = rmp_serde::to_vec_named(&payload)?;
58        let mut encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::fast());
59        encoder.write_all(&msgpack)?;
60        match encoder.finish() {
61            Ok(res) => Ok(res),
62            Err(e) => anyhow::bail!("Error serializing stats payload: {e}"),
63        }
64    }
65
66    pub async fn send_stats_payload<H: HttpClientCapability>(
67        data: Vec<u8>,
68        target: &Endpoint,
69        api_key: &str,
70    ) -> anyhow::Result<()> {
71        let client = H::new_client();
72        let req = http::Request::builder()
73            .method(http::Method::POST)
74            .uri(target.url.clone())
75            .header("Content-Type", "application/msgpack")
76            .header("Content-Encoding", "gzip")
77            .header("DD-API-KEY", api_key)
78            .body(Bytes::from(data))?;
79
80        let response = client
81            .request(req)
82            .await
83            .map_err(|e| anyhow::anyhow!("Failed to send trace stats: {e}"))?;
84
85        if response.status() != http::StatusCode::ACCEPTED {
86            let response_body =
87                String::from_utf8(response.into_body().to_vec()).unwrap_or_default();
88            anyhow::bail!("Server did not accept trace stats: {response_body}");
89        }
90        Ok(())
91    }
92}
93
94#[cfg(test)]
95#[cfg(feature = "mini_agent")]
96mod mini_agent_tests {
97    use crate::stats_utils;
98    use http::Request;
99    use libdd_common::http_common;
100    use libdd_trace_protobuf::pb::{
101        ClientGroupedStats, ClientStatsBucket, ClientStatsPayload, Trilean::NotSet,
102    };
103    use serde_json::Value;
104
105    #[tokio::test]
106    #[cfg_attr(all(miri, target_os = "macos"), ignore)]
107    async fn test_get_stats_from_request_body() {
108        let stats_json = r#"{
109            "Hostname": "TestHost",
110            "Env": "test",
111            "Version": "1.0.0",
112            "Stats": [
113                {
114                    "Start": 0,
115                    "Duration": 10000000000,
116                    "Stats": [
117                        {
118                            "Name": "test-span",
119                            "Service": "test-service",
120                            "Resource": "test-span",
121                            "Type": "",
122                            "HTTPStatusCode": 0,
123                            "Synthetics": false,
124                            "Hits": 1,
125                            "TopLevelHits": 1,
126                            "Errors": 0,
127                            "Duration": 10000000,
128                            "OkSummary": [
129                                0,
130                                0,
131                                0
132                            ],
133                            "ErrorSummary": [
134                                0,
135                                0,
136                                0
137                            ],
138                            "GRPCStatusCode": "0",
139                            "AdditionalMetricTags": [],
140                            "HTTPMethod": "GET",
141                            "HTTPEndpoint": "/test"
142                        }
143                    ]
144                }
145            ],
146            "Lang": "javascript",
147            "TracerVersion": "1.0.0",
148            "RuntimeID": "00000000-0000-0000-0000-000000000000",
149            "Sequence": 1
150        }"#;
151
152        let v: Value = match serde_json::from_str(stats_json) {
153            Ok(value) => value,
154            Err(err) => {
155                panic!("Failed to parse stats JSON: {err}");
156            }
157        };
158
159        let bytes = rmp_serde::to_vec(&v).unwrap();
160        let request = Request::builder()
161            .body(http_common::Body::from(bytes))
162            .unwrap();
163
164        let res = stats_utils::get_stats_from_request_body(request.into_body()).await;
165
166        let client_stats_payload = ClientStatsPayload {
167            hostname: "TestHost".to_string(),
168            env: "test".to_string(),
169            version: "1.0.0".to_string(),
170            stats: vec![ClientStatsBucket {
171                start: 0,
172                duration: 10000000000,
173                stats: vec![ClientGroupedStats {
174                    service: "test-service".to_string(),
175                    name: "test-span".to_string(),
176                    resource: "test-span".to_string(),
177                    http_status_code: 0,
178                    r#type: "".to_string(),
179                    db_type: "".to_string(),
180                    hits: 1,
181                    errors: 0,
182                    duration: 10000000,
183                    ok_summary: vec![0, 0, 0],
184                    error_summary: vec![0, 0, 0],
185                    synthetics: false,
186                    top_level_hits: 1,
187                    span_kind: "".to_string(),
188                    peer_tags: vec![],
189                    is_trace_root: NotSet.into(),
190                    grpc_status_code: "0".to_string(),
191                    http_endpoint: "/test".to_string(),
192                    http_method: "GET".to_string(),
193                    service_source: "".to_string(),
194                    span_derived_primary_tags: vec![],
195                    additional_metric_tags: vec![],
196                }],
197                agent_time_shift: 0,
198            }],
199            lang: "javascript".to_string(),
200            tracer_version: "1.0.0".to_string(),
201            runtime_id: "00000000-0000-0000-0000-000000000000".to_string(),
202            sequence: 1,
203            agent_aggregation: "".to_string(),
204            service: "".to_string(),
205            container_id: "".to_string(),
206            tags: vec![],
207            git_commit_sha: "".to_string(),
208            image_tag: "".to_string(),
209            process_tags_hash: 0,
210            process_tags: "".to_string(),
211        };
212
213        assert!(
214            res.is_ok(),
215            "Expected Ok result, but got Err: {}",
216            res.unwrap_err()
217        );
218        assert_eq!(res.unwrap(), client_stats_payload)
219    }
220
221    #[tokio::test]
222    #[cfg_attr(all(miri, target_os = "macos"), ignore)]
223    async fn test_get_stats_from_request_body_without_stats() {
224        let stats_json = r#"{
225            "Hostname": "TestHost",
226            "Env": "test",
227            "Version": "1.0.0",
228            "Lang": "javascript",
229            "TracerVersion": "1.0.0",
230            "RuntimeID": "00000000-0000-0000-0000-000000000000",
231            "Sequence": 1
232        }"#;
233
234        let v: Value = match serde_json::from_str(stats_json) {
235            Ok(value) => value,
236            Err(err) => {
237                panic!("Failed to parse stats JSON: {err}");
238            }
239        };
240
241        let bytes = rmp_serde::to_vec(&v).unwrap();
242        let request = Request::builder()
243            .body(http_common::Body::from(bytes))
244            .unwrap();
245
246        let res = stats_utils::get_stats_from_request_body(request.into_body()).await;
247
248        let client_stats_payload = ClientStatsPayload {
249            hostname: "TestHost".to_string(),
250            env: "test".to_string(),
251            version: "1.0.0".to_string(),
252            stats: vec![],
253            lang: "javascript".to_string(),
254            tracer_version: "1.0.0".to_string(),
255            runtime_id: "00000000-0000-0000-0000-000000000000".to_string(),
256            sequence: 1,
257            agent_aggregation: "".to_string(),
258            service: "".to_string(),
259            container_id: "".to_string(),
260            tags: vec![],
261            git_commit_sha: "".to_string(),
262            image_tag: "".to_string(),
263            process_tags_hash: 0,
264            process_tags: "".to_string(),
265        };
266
267        assert!(
268            res.is_ok(),
269            "Expected Ok result, but got Err: {}",
270            res.unwrap_err()
271        );
272        assert_eq!(res.unwrap(), client_stats_payload)
273    }
274
275    #[tokio::test]
276    #[cfg_attr(all(miri, target_os = "macos"), ignore)]
277    async fn test_serialize_client_stats_payload_without_stats() {
278        let client_stats_payload_without_stats = ClientStatsPayload {
279            hostname: "TestHost".to_string(),
280            env: "test".to_string(),
281            version: "1.0.0".to_string(),
282            stats: vec![],
283            lang: "javascript".to_string(),
284            tracer_version: "1.0.0".to_string(),
285            runtime_id: "00000000-0000-0000-0000-000000000000".to_string(),
286            sequence: 1,
287            agent_aggregation: "".to_string(),
288            service: "".to_string(),
289            container_id: "".to_string(),
290            tags: vec![],
291            git_commit_sha: "".to_string(),
292            image_tag: "".to_string(),
293            process_tags_hash: 0,
294            process_tags: "".to_string(),
295        };
296
297        let client_stats_payload_without_inner_stats = ClientStatsPayload {
298            hostname: "TestHost".to_string(),
299            env: "test".to_string(),
300            version: "1.0.0".to_string(),
301            stats: vec![ClientStatsBucket {
302                start: 0,
303                duration: 10000000000,
304                stats: vec![],
305                agent_time_shift: 0,
306            }],
307            lang: "javascript".to_string(),
308            tracer_version: "1.0.0".to_string(),
309            runtime_id: "00000000-0000-0000-0000-000000000000".to_string(),
310            sequence: 1,
311            agent_aggregation: "".to_string(),
312            service: "".to_string(),
313            container_id: "".to_string(),
314            tags: vec![],
315            git_commit_sha: "".to_string(),
316            image_tag: "".to_string(),
317            process_tags_hash: 0,
318            process_tags: "".to_string(),
319        };
320
321        let res = stats_utils::serialize_stats_payload(stats_utils::construct_stats_payload(vec![
322            client_stats_payload_without_stats,
323            client_stats_payload_without_inner_stats,
324        ]));
325
326        assert!(
327            res.is_ok(),
328            "Expected Ok result, but got Err: {}",
329            res.unwrap_err()
330        );
331    }
332}
libdd_trace_utils/stats_utils.rs

libdd_trace_utils/
stats_utils.rs
