Skip to content

Commit 4dd532f

Browse files
paullegranddcpaullegranddc
authored
refactor(trace-utils)!: change header name type to accept dynamic values (#1722)
# What does this PR do? * Change the header map type passed throughout data-pipeline and trace utils from `Hashmap<&'static str, String> to `http::HeaderMap` This should not cause extra allocations for fixed header names, as the header names for string values are "const constructed" and trivially copyable. In fact it should cause less allocations as header values are now `http::HeaderValue` instead of `String`. The static ones don't require an allocation and clone becomes a shallow copy. # Motivation OTLP supports requires the ability to defined extra headers sent with the payload in configuration # Additional Notes The first iteration I went through created a `Hashmap<http::HeaderName, String>` but this does not work, as http::HeaderName implement Borrow<str> but does not hash like the &str it represents (see hyperium/http#824)
1 parent 5426a8b commit 4dd532f

File tree

8 files changed

+184
-171
lines changed

8 files changed

+184
-171
lines changed

‎libdd-common/src/lib.rs‎

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -86,10 +86,6 @@ pub mod header {
8686
#![allow(clippy::declare_interior_mutable_const)]
8787
use hyper::{header::HeaderName, http::HeaderValue};
8888

89-
// These strings are defined separately to be used in context where &str are used to represent
90-
// headers (e.g. SendData) while keeping a single source of truth.
91-
pub const DATADOG_SEND_REAL_HTTP_STATUS_STR: &str = "datadog-send-real-http-status";
92-
pub const DATADOG_TRACE_COUNT_STR: &str = "x-datadog-trace-count";
9389
pub const APPLICATION_MSGPACK_STR: &str = "application/msgpack";
9490
pub const APPLICATION_PROTOBUF_STR: &str = "application/x-protobuf";
9591

@@ -101,7 +97,7 @@ pub mod header {
10197
/// If this is not set then the agent will always return a 200 regardless if the payload is
10298
/// dropped.
10399
pub const DATADOG_SEND_REAL_HTTP_STATUS: HeaderName =
104-
HeaderName::from_static(DATADOG_SEND_REAL_HTTP_STATUS_STR);
100+
HeaderName::from_static("datadog-send-real-http-status");
105101
pub const DATADOG_API_KEY: HeaderName = HeaderName::from_static("dd-api-key");
106102
pub const APPLICATION_JSON: HeaderValue = HeaderValue::from_static("application/json");
107103
pub const APPLICATION_MSGPACK: HeaderValue = HeaderValue::from_static(APPLICATION_MSGPACK_STR);

‎libdd-data-pipeline/src/stats_exporter.rs‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@
33

44
use std::{
55
borrow::Borrow,
6-
collections::HashMap,
76
sync::{
87
atomic::{AtomicU64, Ordering},
98
Arc, Mutex,
@@ -84,11 +83,11 @@ impl StatsExporter {
8483
}
8584
let body = rmp_serde::encode::to_vec_named(&payload)?;
8685

87-
let mut headers: HashMap<&'static str, String> = self.meta.borrow().into();
86+
let mut headers: http::HeaderMap = self.meta.borrow().into();
8887

8988
headers.insert(
90-
http::header::CONTENT_TYPE.as_str(),
91-
libdd_common::header::APPLICATION_MSGPACK_STR.to_string(),
89+
http::header::CONTENT_TYPE,
90+
libdd_common::header::APPLICATION_MSGPACK,
9291
);
9392

9493
let result = send_with_retry(

‎libdd-data-pipeline/src/trace_exporter/agent_response.rs‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,13 @@
11
// Copyright 2024-Present Datadog, Inc. https://www.datadoghq.com/
22
// SPDX-License-Identifier: Apache-2.0
33

4+
use http::HeaderName;
45
use std::sync::Arc;
56

67
use arc_swap::ArcSwap;
78

8-
pub const DATADOG_RATES_PAYLOAD_VERSION_HEADER: &str = "datadog-rates-payload-version";
9+
pub const DATADOG_RATES_PAYLOAD_VERSION: HeaderName =
10+
HeaderName::from_static("datadog-rates-payload-version");
911

1012
/// `AgentResponse` structure holds agent response information upon successful request.
1113
#[derive(Debug, PartialEq)]

‎libdd-data-pipeline/src/trace_exporter/mod.rs‎

Lines changed: 7 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ use crate::pausable_worker::PausableWorker;
1919
use crate::stats_exporter::StatsExporter;
2020
use crate::telemetry::{SendPayloadTelemetry, TelemetryClient};
2121
use crate::trace_exporter::agent_response::{
22-
AgentResponsePayloadVersion, DATADOG_RATES_PAYLOAD_VERSION_HEADER,
22+
AgentResponsePayloadVersion, DATADOG_RATES_PAYLOAD_VERSION,
2323
};
2424
use crate::trace_exporter::error::{InternalErrorKind, RequestError, TraceExporterError};
2525
use crate::{
@@ -45,7 +45,7 @@ use libdd_trace_utils::trace_utils::TracerHeaderTags;
4545
use std::io;
4646
use std::sync::{Arc, Mutex};
4747
use std::time::Duration;
48-
use std::{borrow::Borrow, collections::HashMap, str::FromStr};
48+
use std::{borrow::Borrow, str::FromStr};
4949
use tokio::runtime::Runtime;
5050
use tracing::{debug, error, warn};
5151

@@ -148,8 +148,8 @@ impl<'a> From<&'a TracerMetadata> for TracerHeaderTags<'a> {
148148
}
149149
}
150150

151-
impl<'a> From<&'a TracerMetadata> for HashMap<&'static str, String> {
152-
fn from(tags: &'a TracerMetadata) -> HashMap<&'static str, String> {
151+
impl<'a> From<&'a TracerMetadata> for http::HeaderMap {
152+
fn from(tags: &'a TracerMetadata) -> http::HeaderMap {
153153
TracerHeaderTags::from(tags).into()
154154
}
155155
}
@@ -563,7 +563,7 @@ impl TraceExporter {
563563
&self,
564564
endpoint: &Endpoint,
565565
mp_payload: Vec<u8>,
566-
headers: HashMap<&'static str, String>,
566+
headers: http::HeaderMap,
567567
chunks: usize,
568568
chunks_dropped_p0: usize,
569569
) -> Result<AgentResponse, TraceExporterError> {
@@ -738,7 +738,7 @@ impl TraceExporter {
738738
match (
739739
status.is_success(),
740740
self.agent_payload_response_version.as_ref(),
741-
response.headers().get(DATADOG_RATES_PAYLOAD_VERSION_HEADER),
741+
response.headers().get(DATADOG_RATES_PAYLOAD_VERSION),
742742
) {
743743
(false, _, _) => {
744744
// If the status is not success, the rates are considered unchanged
@@ -876,7 +876,6 @@ mod tests {
876876
use libdd_trace_utils::msgpack_encoder;
877877
use libdd_trace_utils::span::v04::SpanBytes;
878878
use libdd_trace_utils::span::v05;
879-
use std::collections::HashMap;
880879
use std::net;
881880
use std::time::Duration;
882881
use tokio::time::sleep;
@@ -920,7 +919,7 @@ mod tests {
920919
..Default::default()
921920
};
922921

923-
let hashmap: HashMap<&'static str, String> = (&tracer_tags).into();
922+
let hashmap: http::HeaderMap = (&tracer_tags).into();
924923

925924
assert_eq!(hashmap.get("datadog-meta-tracer-version").unwrap(), "v0.1");
926925
assert_eq!(hashmap.get("datadog-meta-lang").unwrap(), "rust");

‎libdd-data-pipeline/src/trace_exporter/trace_serializer.rs‎

Lines changed: 22 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -2,27 +2,26 @@
22
// SPDX-License-Identifier: Apache-2.0
33

44
use crate::trace_exporter::agent_response::{
5-
AgentResponsePayloadVersion, DATADOG_RATES_PAYLOAD_VERSION_HEADER,
5+
AgentResponsePayloadVersion, DATADOG_RATES_PAYLOAD_VERSION,
66
};
77
use crate::trace_exporter::error::TraceExporterError;
88
use crate::trace_exporter::TraceExporterOutputFormat;
9-
use http::header::CONTENT_TYPE;
9+
use http::{header::CONTENT_TYPE, HeaderMap, HeaderValue};
1010
use libdd_common::header::{
11-
APPLICATION_MSGPACK_STR, DATADOG_SEND_REAL_HTTP_STATUS_STR, DATADOG_TRACE_COUNT_STR,
11+
APPLICATION_MSGPACK, DATADOG_SEND_REAL_HTTP_STATUS, DATADOG_TRACE_COUNT,
1212
};
1313
use libdd_trace_utils::msgpack_decoder::decode::error::DecodeError;
1414
use libdd_trace_utils::msgpack_encoder;
1515
use libdd_trace_utils::span::{v04::Span, TraceData};
1616
use libdd_trace_utils::trace_utils::{self, TracerHeaderTags};
1717
use libdd_trace_utils::tracer_payload;
18-
use std::collections::HashMap;
1918

2019
/// Prepared traces payload ready for sending to the agent
2120
pub(super) struct PreparedTracesPayload {
2221
/// Serialized msgpack payload
2322
pub data: Vec<u8>,
2423
/// HTTP headers for the request
25-
pub headers: HashMap<&'static str, String>,
24+
pub headers: HeaderMap,
2625
/// Number of trace chunks
2726
pub chunk_count: usize,
2827
}
@@ -78,20 +77,16 @@ impl<'a> TraceSerializer<'a> {
7877
}
7978

8079
/// Build HTTP headers for traces request
81-
fn build_traces_headers(
82-
&self,
83-
header_tags: TracerHeaderTags,
84-
chunk_count: usize,
85-
) -> HashMap<&'static str, String> {
86-
let mut headers: HashMap<&'static str, String> = header_tags.into();
87-
headers.insert(DATADOG_SEND_REAL_HTTP_STATUS_STR, "1".to_string());
88-
headers.insert(DATADOG_TRACE_COUNT_STR, chunk_count.to_string());
89-
headers.insert(CONTENT_TYPE.as_str(), APPLICATION_MSGPACK_STR.to_string());
80+
fn build_traces_headers(&self, header_tags: TracerHeaderTags, chunk_count: usize) -> HeaderMap {
81+
let mut headers: HeaderMap = header_tags.into();
82+
headers.reserve(4);
83+
headers.insert(DATADOG_SEND_REAL_HTTP_STATUS, HeaderValue::from_static("1"));
84+
headers.insert(DATADOG_TRACE_COUNT, chunk_count.into());
85+
headers.insert(CONTENT_TYPE, APPLICATION_MSGPACK);
9086
if let Some(agent_payload_response_version) = &self.agent_payload_response_version {
91-
headers.insert(
92-
DATADOG_RATES_PAYLOAD_VERSION_HEADER,
93-
agent_payload_response_version.header_value(),
94-
);
87+
// should never fail, as the version should only contain visible ascii chars
88+
let _ = HeaderValue::try_from(agent_payload_response_version.header_value())
89+
.map(|v| headers.insert(DATADOG_RATES_PAYLOAD_VERSION, v));
9590
}
9691
headers
9792
}
@@ -115,9 +110,7 @@ mod tests {
115110
use super::*;
116111
use crate::trace_exporter::agent_response::AgentResponsePayloadVersion;
117112
use http::header::CONTENT_TYPE;
118-
use libdd_common::header::{
119-
APPLICATION_MSGPACK_STR, DATADOG_SEND_REAL_HTTP_STATUS_STR, DATADOG_TRACE_COUNT_STR,
120-
};
113+
use libdd_common::header::APPLICATION_MSGPACK_STR;
121114
use libdd_tinybytes::BytesString;
122115
use libdd_trace_utils::span::v04::SpanBytes;
123116
use libdd_trace_utils::trace_utils::TracerHeaderTags;
@@ -179,12 +172,9 @@ mod tests {
179172
let headers = serializer.build_traces_headers(header_tags, 3);
180173

181174
// Check basic headers are present
182-
assert_eq!(headers.get(DATADOG_SEND_REAL_HTTP_STATUS_STR).unwrap(), "1");
183-
assert_eq!(headers.get(DATADOG_TRACE_COUNT_STR).unwrap(), "3");
184-
assert_eq!(
185-
headers.get(CONTENT_TYPE.as_str()).unwrap(),
186-
APPLICATION_MSGPACK_STR
187-
);
175+
assert_eq!(headers.get(DATADOG_SEND_REAL_HTTP_STATUS).unwrap(), "1");
176+
assert_eq!(headers.get(DATADOG_TRACE_COUNT).unwrap(), "3");
177+
assert_eq!(headers.get(CONTENT_TYPE).unwrap(), APPLICATION_MSGPACK_STR);
188178

189179
// Check tracer metadata headers are present
190180
assert_eq!(headers.get("datadog-meta-lang").unwrap(), "rust");
@@ -212,8 +202,8 @@ mod tests {
212202
let headers = serializer.build_traces_headers(header_tags, 2);
213203

214204
// Check that agent payload version header is included
215-
assert!(headers.contains_key(DATADOG_RATES_PAYLOAD_VERSION_HEADER));
216-
assert_eq!(headers.get(DATADOG_TRACE_COUNT_STR).unwrap(), "2");
205+
assert!(headers.contains_key(DATADOG_RATES_PAYLOAD_VERSION));
206+
assert_eq!(headers.get(DATADOG_TRACE_COUNT).unwrap(), "2");
217207
}
218208

219209
#[test]
@@ -346,7 +336,7 @@ mod tests {
346336
assert!(!prepared.headers.is_empty());
347337

348338
// Check headers
349-
assert_eq!(prepared.headers.get(DATADOG_TRACE_COUNT_STR).unwrap(), "2");
339+
assert_eq!(prepared.headers.get(DATADOG_TRACE_COUNT).unwrap(), "2");
350340
assert_eq!(prepared.headers.get("datadog-meta-lang").unwrap(), "rust");
351341
}
352342

@@ -377,9 +367,7 @@ mod tests {
377367

378368
let prepared = result.unwrap();
379369
assert_eq!(prepared.chunk_count, 1);
380-
assert!(prepared
381-
.headers
382-
.contains_key(DATADOG_RATES_PAYLOAD_VERSION_HEADER));
370+
assert!(prepared.headers.contains_key(DATADOG_RATES_PAYLOAD_VERSION));
383371
}
384372

385373
#[test]
@@ -394,7 +382,7 @@ mod tests {
394382
let prepared = result.unwrap();
395383
assert_eq!(prepared.chunk_count, 0);
396384
assert!(!prepared.data.is_empty()); // Even empty traces result in some serialized data
397-
assert_eq!(prepared.headers.get(DATADOG_TRACE_COUNT_STR).unwrap(), "0");
385+
assert_eq!(prepared.headers.get(DATADOG_TRACE_COUNT).unwrap(), "0");
398386
}
399387

400388
#[test]

0 commit comments

Comments
 (0)