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                            "HTTPMethod": "GET",
140                            "HTTPEndpoint": "/test"
141                        }
142                    ]
143                }
144            ],
145            "Lang": "javascript",
146            "TracerVersion": "1.0.0",
147            "RuntimeID": "00000000-0000-0000-0000-000000000000",
148            "Sequence": 1
149        }"#;
150
151        let v: Value = match serde_json::from_str(stats_json) {
152            Ok(value) => value,
153            Err(err) => {
154                panic!("Failed to parse stats JSON: {err}");
155            }
156        };
157
158        let bytes = rmp_serde::to_vec(&v).unwrap();
159        let request = Request::builder()
160            .body(http_common::Body::from(bytes))
161            .unwrap();
162
163        let res = stats_utils::get_stats_from_request_body(request.into_body()).await;
164
165        let client_stats_payload = ClientStatsPayload {
166            hostname: "TestHost".to_string(),
167            env: "test".to_string(),
168            version: "1.0.0".to_string(),
169            stats: vec![ClientStatsBucket {
170                start: 0,
171                duration: 10000000000,
172                stats: vec![ClientGroupedStats {
173                    service: "test-service".to_string(),
174                    name: "test-span".to_string(),
175                    resource: "test-span".to_string(),
176                    http_status_code: 0,
177                    r#type: "".to_string(),
178                    db_type: "".to_string(),
179                    hits: 1,
180                    errors: 0,
181                    duration: 10000000,
182                    ok_summary: vec![0, 0, 0],
183                    error_summary: vec![0, 0, 0],
184                    synthetics: false,
185                    top_level_hits: 1,
186                    span_kind: "".to_string(),
187                    peer_tags: vec![],
188                    is_trace_root: NotSet.into(),
189                    grpc_status_code: "0".to_string(),
190                    http_endpoint: "/test".to_string(),
191                    http_method: "GET".to_string(),
192                    service_source: "".to_string(),
193                    span_derived_primary_tags: vec![],
194                    additional_metric_tags: vec![],
195                }],
196                agent_time_shift: 0,
197            }],
198            lang: "javascript".to_string(),
199            tracer_version: "1.0.0".to_string(),
200            runtime_id: "00000000-0000-0000-0000-000000000000".to_string(),
201            sequence: 1,
202            agent_aggregation: "".to_string(),
203            service: "".to_string(),
204            container_id: "".to_string(),
205            tags: vec![],
206            git_commit_sha: "".to_string(),
207            image_tag: "".to_string(),
208            process_tags_hash: 0,
209            process_tags: "".to_string(),
210        };
211
212        assert!(
213            res.is_ok(),
214            "Expected Ok result, but got Err: {}",
215            res.unwrap_err()
216        );
217        assert_eq!(res.unwrap(), client_stats_payload)
218    }
219
220    #[tokio::test]
221    #[cfg_attr(all(miri, target_os = "macos"), ignore)]
222    async fn test_get_stats_from_request_body_without_stats() {
223        let stats_json = r#"{
224            "Hostname": "TestHost",
225            "Env": "test",
226            "Version": "1.0.0",
227            "Lang": "javascript",
228            "TracerVersion": "1.0.0",
229            "RuntimeID": "00000000-0000-0000-0000-000000000000",
230            "Sequence": 1
231        }"#;
232
233        let v: Value = match serde_json::from_str(stats_json) {
234            Ok(value) => value,
235            Err(err) => {
236                panic!("Failed to parse stats JSON: {err}");
237            }
238        };
239
240        let bytes = rmp_serde::to_vec(&v).unwrap();
241        let request = Request::builder()
242            .body(http_common::Body::from(bytes))
243            .unwrap();
244
245        let res = stats_utils::get_stats_from_request_body(request.into_body()).await;
246
247        let client_stats_payload = ClientStatsPayload {
248            hostname: "TestHost".to_string(),
249            env: "test".to_string(),
250            version: "1.0.0".to_string(),
251            stats: vec![],
252            lang: "javascript".to_string(),
253            tracer_version: "1.0.0".to_string(),
254            runtime_id: "00000000-0000-0000-0000-000000000000".to_string(),
255            sequence: 1,
256            agent_aggregation: "".to_string(),
257            service: "".to_string(),
258            container_id: "".to_string(),
259            tags: vec![],
260            git_commit_sha: "".to_string(),
261            image_tag: "".to_string(),
262            process_tags_hash: 0,
263            process_tags: "".to_string(),
264        };
265
266        assert!(
267            res.is_ok(),
268            "Expected Ok result, but got Err: {}",
269            res.unwrap_err()
270        );
271        assert_eq!(res.unwrap(), client_stats_payload)
272    }
273
274    #[tokio::test]
275    #[cfg_attr(all(miri, target_os = "macos"), ignore)]
276    async fn test_serialize_client_stats_payload_without_stats() {
277        let client_stats_payload_without_stats = ClientStatsPayload {
278            hostname: "TestHost".to_string(),
279            env: "test".to_string(),
280            version: "1.0.0".to_string(),
281            stats: vec![],
282            lang: "javascript".to_string(),
283            tracer_version: "1.0.0".to_string(),
284            runtime_id: "00000000-0000-0000-0000-000000000000".to_string(),
285            sequence: 1,
286            agent_aggregation: "".to_string(),
287            service: "".to_string(),
288            container_id: "".to_string(),
289            tags: vec![],
290            git_commit_sha: "".to_string(),
291            image_tag: "".to_string(),
292            process_tags_hash: 0,
293            process_tags: "".to_string(),
294        };
295
296        let client_stats_payload_without_inner_stats = ClientStatsPayload {
297            hostname: "TestHost".to_string(),
298            env: "test".to_string(),
299            version: "1.0.0".to_string(),
300            stats: vec![ClientStatsBucket {
301                start: 0,
302                duration: 10000000000,
303                stats: vec![],
304                agent_time_shift: 0,
305            }],
306            lang: "javascript".to_string(),
307            tracer_version: "1.0.0".to_string(),
308            runtime_id: "00000000-0000-0000-0000-000000000000".to_string(),
309            sequence: 1,
310            agent_aggregation: "".to_string(),
311            service: "".to_string(),
312            container_id: "".to_string(),
313            tags: vec![],
314            git_commit_sha: "".to_string(),
315            image_tag: "".to_string(),
316            process_tags_hash: 0,
317            process_tags: "".to_string(),
318        };
319
320        let res = stats_utils::serialize_stats_payload(stats_utils::construct_stats_payload(vec![
321            client_stats_payload_without_stats,
322            client_stats_payload_without_inner_stats,
323        ]));
324
325        assert!(
326            res.is_ok(),
327            "Expected Ok result, but got Err: {}",
328            res.unwrap_err()
329        );
330    }
331}
libdd_trace_utils/stats_utils.rs

libdd_trace_utils/
stats_utils.rs
