Skip to content

Commit 6c2a87d

Browse files
danielsndanielsn
committed
add a reverse map from sample types to indicies
1 parent 4a39eaa commit 6c2a87d

5 files changed

Lines changed: 86 additions & 52 deletions

File tree

‎Cargo.lock‎

Lines changed: 21 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎libdd-profiling/Cargo.toml‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ bench = false
1919
[features]
2020
default = []
2121
cxx = ["dep:cxx", "dep:cxx-build"]
22-
otel = ["dep:opentelemetry-proto"]
22+
otel = ["dep:opentelemetry-proto", "dep:enum-map"]
2323
test-utils = ["libdd-common/test-utils"]
2424

2525
[[bench]]
@@ -46,6 +46,7 @@ libdd-alloc = { version = "1.0.0", path = "../libdd-alloc" }
4646
libdd-common = { version = "1.1.0", path = "../libdd-common", default-features = false, features = ["reqwest", "test-utils"] }
4747
libdd-profiling-protobuf = { version = "1.0.0", path = "../libdd-profiling-protobuf", features = ["prost_impls"] }
4848
mime = "0.3.16"
49+
enum-map = { version = "2", optional = true }
4950
opentelemetry-proto = { version = "0.31", optional = true, default-features = false, features = ["profiles", "gen-tonic-messages"] }
5051
parking_lot = { version = "0.12", default-features = false }
5152
prost = "0.14.1"

‎libdd-profiling/src/api/sample_type.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ use super::ValueType;
1212
/// - **dd-trace-dotnet**: Sample type definitions for allocations, locks, CPU, walltime,
1313
/// exceptions, live objects, HTTP requests
1414
/// - **pprof-nodejs**: `profile-serializer.ts` (value type functions)
15+
#[cfg_attr(feature = "otel", derive(enum_map::Enum))]
1516
#[cfg_attr(test, derive(bolero::generator::TypeGenerator, strum::EnumIter))]
1617
#[repr(C)]
1718
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]

‎libdd-profiling/src/internal/otel_style_observation.rs‎

Lines changed: 16 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -5,12 +5,15 @@ use crate::{
55
api::SampleType,
66
internal::{Sample, Timestamp},
77
};
8+
use enum_map::EnumMap;
89
use std::collections::HashMap;
10+
use std::sync::Arc;
911

1012
#[derive(Default)]
1113
#[allow(clippy::type_complexity)]
1214
pub struct Observations {
13-
sample_types: Box<[SampleType]>,
15+
sample_types: Arc<[SampleType]>,
16+
sample_type_index_map: EnumMap<SampleType, Option<usize>>,
1417
pub aggregated: HashMap<SampleType, HashMap<Sample, i64>>,
1518
pub aggregated2: HashMap<(SampleType, SampleType), HashMap<Sample, (i64, i64)>>,
1619
pub timestamped: HashMap<SampleType, HashMap<Sample, Vec<(i64, Timestamp)>>>,
@@ -19,9 +22,13 @@ pub struct Observations {
1922
}
2023

2124
impl Observations {
22-
pub fn new(sample_types: Box<[SampleType]>) -> Self {
25+
pub fn new(
26+
sample_types: Arc<[SampleType]>,
27+
sample_type_index_map: EnumMap<SampleType, Option<usize>>,
28+
) -> Self {
2329
Self {
2430
sample_types,
31+
sample_type_index_map,
2532
..Default::default()
2633
}
2734
}
@@ -130,26 +137,18 @@ impl Observations {
130137
// TODO make this a trait impl
131138
#[allow(clippy::should_implement_trait)]
132139
pub fn into_iter(self) -> impl Iterator<Item = (Sample, Option<Timestamp>, Vec<i64>)> {
133-
let index_map: HashMap<SampleType, usize> = self
134-
.sample_types
135-
.iter()
136-
.enumerate()
137-
.map(|(idx, typ)| (*typ, idx))
138-
.collect();
140+
let index_map = self.sample_type_index_map;
139141

140142
// Invariant: all keys in aggregated/aggregated2/timestamped/timestamped2 come from
141143
// add(), which indexes into self.sample_types for non-zero values. So every
142144
// sample_type key is present in index_map.
143145
let len: usize = self.sample_types.len();
144-
let index_map_ts = index_map.clone();
145-
let index_map_ts2 = index_map.clone();
146-
let index_map_accum2 = index_map.clone();
147146
let accum_iter = self
148147
.aggregated
149148
.into_iter()
150149
.flat_map(move |(sample_type, inner)| {
151150
#[allow(clippy::unwrap_used)]
152-
let index = *index_map.get(&sample_type).unwrap();
151+
let index = index_map[sample_type].unwrap();
153152
inner.into_iter().map(move |(sample, value)| {
154153
let mut vals = vec![0; len];
155154
vals[index] = value;
@@ -162,9 +161,9 @@ impl Observations {
162161
.into_iter()
163162
.flat_map(move |((st1, st2), inner)| {
164163
#[allow(clippy::unwrap_used)]
165-
let index1 = *index_map_accum2.get(&st1).unwrap();
164+
let index1 = index_map[st1].unwrap();
166165
#[allow(clippy::unwrap_used)]
167-
let index2 = *index_map_accum2.get(&st2).unwrap();
166+
let index2 = index_map[st2].unwrap();
168167
inner.into_iter().map(move |(sample, (val1, val2))| {
169168
let mut vals = vec![0; len];
170169
vals[index1] = val1;
@@ -173,24 +172,13 @@ impl Observations {
173172
})
174173
});
175174

176-
// let mut accum: HashMap<Sample, Vec<i64>> = HashMap::new();
177-
// for (sample_type, samples) in self.aggregated.into_iter() {
178-
// let Some(idx) = index_map.get(&sample_type) else {
179-
// continue;
180-
// };
181-
// for (sample, val) in samples.into_iter() {
182-
// let val_accum = accum.entry(sample).or_insert_with(|| vec![0; len]);
183-
// val_accum[*idx] += val;
184-
// }
185-
// }
186-
// let accum_iter = accum.into_iter().map(|(k, v)| (k, None, v));
187175

188176
let ts_iter = self
189177
.timestamped
190178
.into_iter()
191179
.flat_map(move |(sample_type, inner)| {
192180
#[allow(clippy::unwrap_used)]
193-
let index = *index_map_ts.get(&sample_type).unwrap();
181+
let index = index_map[sample_type].unwrap();
194182
inner.into_iter().flat_map(move |(sample, ts_vals)| {
195183
ts_vals.into_iter().map(move |(value, ts)| {
196184
let mut vals = vec![0; len];
@@ -205,9 +193,9 @@ impl Observations {
205193
.into_iter()
206194
.flat_map(move |((st1, st2), inner)| {
207195
#[allow(clippy::unwrap_used)]
208-
let index1 = *index_map_ts2.get(&st1).unwrap();
196+
let index1 = index_map[st1].unwrap();
209197
#[allow(clippy::unwrap_used)]
210-
let index2 = *index_map_ts2.get(&st2).unwrap();
198+
let index2 = index_map[st2].unwrap();
211199
inner.into_iter().flat_map(move |(sample, ts_vals)| {
212200
ts_vals.into_iter().map(move |(val1, val2, ts)| {
213201
let mut vals = vec![0; len];

‎libdd-profiling/src/internal/profile/mod.rs‎

Lines changed: 46 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,15 @@ use std::sync::atomic::AtomicU64;
3131
use std::sync::{Arc, Mutex};
3232
use std::time::{Duration, SystemTime};
3333

34+
#[cfg(feature = "otel")]
35+
use enum_map::EnumMap;
36+
#[cfg(feature = "otel")]
37+
use opentelemetry_proto::tonic::common::v1::{any_value, AnyValue};
38+
#[cfg(feature = "otel")]
39+
use prost::Message;
40+
#[cfg(feature = "otel")]
41+
use std::io::Write;
42+
3443
pub struct Profile {
3544
/// Translates from the new Id2 APIs to the older internal APIs. Long-term,
3645
/// this should probably use the dictionary directly.
@@ -45,7 +54,9 @@ pub struct Profile {
4554
mappings: FxIndexSet<Mapping>,
4655
observations: Observations,
4756
period: Option<api::Period>,
48-
sample_types: Box<[api::SampleType]>,
57+
#[cfg(feature = "otel")]
58+
sample_type_index_map: EnumMap<api::SampleType, Option<usize>>,
59+
sample_types: std::sync::Arc<[api::SampleType]>,
4960
stack_traces: FxIndexSet<StackTrace>,
5061
start_time: SystemTime,
5162
strings: StringTable,
@@ -391,7 +402,7 @@ impl Profile {
391402
sample_types: &[api::SampleType],
392403
period: Option<api::Period>,
393404
) -> io::Result<Self> {
394-
Self::try_new_internal(period, sample_types.to_vec().into_boxed_slice(), None, None)
405+
Self::try_new_internal(period, std::sync::Arc::from(sample_types), None, None)
395406
}
396407

397408
/// Tries to create a profile with the given period and sample types.
@@ -404,7 +415,7 @@ impl Profile {
404415
) -> io::Result<Self> {
405416
Self::try_new_internal(
406417
period,
407-
sample_types.to_vec().into_boxed_slice(),
418+
std::sync::Arc::from(sample_types),
408419
None,
409420
Some(ProfilesDictionaryTranslator::new(profiles_dictionary)),
410421
)
@@ -418,7 +429,7 @@ impl Profile {
418429
) -> io::Result<Self> {
419430
Self::try_new_internal(
420431
period,
421-
sample_types.to_vec().into_boxed_slice(),
432+
std::sync::Arc::from(sample_types),
422433
Some(string_storage),
423434
None,
424435
)
@@ -445,7 +456,7 @@ impl Profile {
445456

446457
let mut profile = Profile::try_new_internal(
447458
self.period,
448-
self.sample_types.clone(),
459+
std::sync::Arc::clone(&self.sample_types),
449460
self.string_storage.clone(),
450461
profiles_dictionary_translator,
451462
)
@@ -559,7 +570,7 @@ impl Profile {
559570
let period = self.period.map(|period| period.value).unwrap_or(0);
560571

561572
let sample_types = self.sample_types.clone();
562-
let mut otel_profiles: HashMap<SampleType, crate::otel::Profile> = sample_types
573+
let mut otel_profiles: HashMap<api::SampleType, crate::otel::Profile> = sample_types
563574
.iter()
564575
.map(|st| {
565576
let sample_type = self.intern_sample_type(*st);
@@ -651,13 +662,6 @@ impl Profile {
651662
})
652663
.collect();
653664

654-
use std::io::Write;
655-
656-
use opentelemetry_proto::tonic::common::v1::{any_value, AnyValue};
657-
use prost::Message;
658-
659-
use crate::api::SampleType;
660-
661665
let attribute_table: Vec<crate::otel::KeyValueAndUnit> =
662666
std::iter::once(crate::otel::KeyValueAndUnit::default())
663667
.chain(std::mem::take(&mut self.labels).into_iter().map(|label| {
@@ -1283,10 +1287,17 @@ impl Profile {
12831287
/// the owned values.
12841288
fn try_new_internal(
12851289
period: Option<api::Period>,
1286-
sample_types: Box<[api::SampleType]>,
1290+
sample_types: std::sync::Arc<[api::SampleType]>,
12871291
string_storage: Option<Arc<Mutex<ManagedStringStorage>>>,
12881292
profiles_dictionary_translator: Option<ProfilesDictionaryTranslator>,
12891293
) -> io::Result<Self> {
1294+
#[cfg(feature = "otel")]
1295+
let sample_type_index_map = sample_types
1296+
.iter()
1297+
.enumerate()
1298+
.map(|(idx, &st)| (st, Some(idx)))
1299+
.collect::<EnumMap<api::SampleType, Option<usize>>>();
1300+
12901301
let start_time = SystemTime::now();
12911302
let mut profile = Self {
12921303
profiles_dictionary_translator,
@@ -1301,6 +1312,8 @@ impl Profile {
13011312
mappings: Default::default(),
13021313
observations: Default::default(),
13031314
period,
1315+
#[cfg(feature = "otel")]
1316+
sample_type_index_map,
13041317
sample_types,
13051318
stack_traces: Default::default(),
13061319
start_time,
@@ -1339,7 +1352,10 @@ impl Profile {
13391352
profile.observations = {
13401353
#[cfg(feature = "otel")]
13411354
{
1342-
Observations::new(profile.sample_types.clone())
1355+
Observations::new(
1356+
std::sync::Arc::clone(&profile.sample_types),
1357+
profile.sample_type_index_map,
1358+
)
13431359
}
13441360
#[cfg(not(feature = "otel"))]
13451361
{
@@ -1616,10 +1632,7 @@ mod api_tests {
16161632

16171633
let dict = data.dictionary.as_ref().expect("dictionary");
16181634
let scope = &data.resource_profiles[0].scope_profiles[0];
1619-
let otel_profile = scope
1620-
.profiles
1621-
.first()
1622-
.expect("one profile for CpuSamples");
1635+
let otel_profile = scope.profiles.first().expect("one profile for CpuSamples");
16231636

16241637
assert_eq!(otel_profile.sample.len(), 3);
16251638
assert_eq!(dict.mapping_table.len(), 2); // default + 1 real
@@ -1639,7 +1652,8 @@ mod api_tests {
16391652
assert_eq!(value, 101);
16401653
}
16411654

1642-
// OTEL stores timestamps in timestamps_unix_nano, not as attributes (unlike pprof's end_timestamp_ns label)
1655+
// OTEL stores timestamps in timestamps_unix_nano, not as attributes (unlike pprof's
1656+
// end_timestamp_ns label)
16431657
let sample_with_timestamp = samples
16441658
.iter()
16451659
.find(|s| !s.timestamps_unix_nano.is_empty())
@@ -1714,7 +1728,10 @@ mod api_tests {
17141728
.iter()
17151729
.find(|s| otel_sample_attr_int(dict, s, "local root span id") == Some(10))
17161730
.expect("sample with span id 10");
1717-
assert_eq!(otel_sample_attr_int(dict, s1, "local root span id"), Some(10));
1731+
assert_eq!(
1732+
otel_sample_attr_int(dict, s1, "local root span id"),
1733+
Some(10)
1734+
);
17181735
assert_eq!(otel_sample_attr_str(dict, s1, "other"), Some("test"));
17191736
assert_eq!(
17201737
otel_sample_attr_str(dict, s1, "trace endpoint"),
@@ -1725,7 +1742,10 @@ mod api_tests {
17251742
.iter()
17261743
.find(|s| otel_sample_attr_int(dict, s, "local root span id") == Some(11))
17271744
.expect("sample with span id 11");
1728-
assert_eq!(otel_sample_attr_int(dict, s2, "local root span id"), Some(11));
1745+
assert_eq!(
1746+
otel_sample_attr_int(dict, s2, "local root span id"),
1747+
Some(11)
1748+
);
17291749
assert_eq!(otel_sample_attr_str(dict, s2, "other"), Some("test"));
17301750
assert_eq!(otel_sample_attr_str(dict, s2, "trace endpoint"), None);
17311751
Ok(())
@@ -1798,7 +1818,10 @@ mod api_tests {
17981818
.iter()
17991819
.find(|s| otel_sample_attr_int(dict, s, "local root span id") == Some(10))
18001820
.expect("sample with span id 10");
1801-
assert_eq!(otel_sample_attr_int(dict, s1, "local root span id"), Some(10));
1821+
assert_eq!(
1822+
otel_sample_attr_int(dict, s1, "local root span id"),
1823+
Some(10)
1824+
);
18021825
assert_eq!(
18031826
otel_sample_attr_str(dict, s1, "trace endpoint"),
18041827
Some("endpoint 10")

0 commit comments

Comments
 (0)