libdd_profiling/internal/profile/interning_api/
mod.rs

1// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/
2// SPDX-License-Identifier: Apache-2.0
3
4mod generational_ids;
5pub use generational_ids::*;
6
7use crate::api::ManagedStringId;
8use crate::collections::identifiable::{Dedup, StringId};
9use crate::internal::{
10    Function, FunctionId, Label, LabelId, LabelSet, LabelSetId, Location, LocationId, Mapping,
11    MappingId, Profile, Sample, StackTrace, StackTraceId, Timestamp,
12};
13use std::sync::atomic::Ordering::SeqCst;
14
15impl Profile {
16    pub fn intern_function(
17        &mut self,
18        name: GenerationalId<StringId>,
19        system_name: GenerationalId<StringId>,
20        filename: GenerationalId<StringId>,
21    ) -> anyhow::Result<GenerationalId<FunctionId>> {
22        let function = Function {
23            name: name.get(self.generation)?,
24            system_name: system_name.get(self.generation)?,
25            filename: filename.get(self.generation)?,
26        };
27        let id = self.functions.dedup(function);
28        Ok(GenerationalId::new(id, self.generation))
29    }
30
31    pub fn intern_label_num(
32        &mut self,
33        key: GenerationalId<StringId>,
34        val: i64,
35        unit: GenerationalId<StringId>,
36    ) -> anyhow::Result<GenerationalId<LabelId>> {
37        let key = key.get(self.generation)?;
38        let unit = unit.get(self.generation)?;
39        let id = self.labels.dedup(Label::num(key, val, unit));
40        Ok(GenerationalId::new(id, self.generation))
41    }
42
43    pub fn intern_label_str(
44        &mut self,
45        key: GenerationalId<StringId>,
46        val: GenerationalId<StringId>,
47    ) -> anyhow::Result<GenerationalId<LabelId>> {
48        let key = key.get(self.generation)?;
49        let val = val.get(self.generation)?;
50        let id = self.labels.dedup(Label::str(key, val));
51        Ok(GenerationalId::new(id, self.generation))
52    }
53
54    pub fn intern_labelset(
55        &mut self,
56        labels: &[GenerationalId<LabelId>],
57    ) -> anyhow::Result<GenerationalId<LabelSetId>> {
58        let labels = labels
59            .iter()
60            .map(|l| l.get(self.generation))
61            .collect::<anyhow::Result<Box<_>>>()?;
62        let labels = LabelSet::new(labels);
63        let id = self.label_sets.dedup(labels);
64        Ok(GenerationalId::new(id, self.generation))
65    }
66
67    pub fn intern_location(
68        &mut self,
69        mapping_id: Option<GenerationalId<MappingId>>,
70        function_id: GenerationalId<FunctionId>,
71        address: u64,
72        line: i64,
73    ) -> anyhow::Result<GenerationalId<LocationId>> {
74        let location = Location {
75            mapping_id: mapping_id.map(|id| id.get(self.generation)).transpose()?,
76            function_id: function_id.get(self.generation)?,
77            address,
78            line,
79        };
80        let id = self.locations.dedup(location);
81        Ok(GenerationalId::new(id, self.generation))
82    }
83
84    pub fn intern_managed_string(
85        &mut self,
86        s: ManagedStringId,
87    ) -> anyhow::Result<GenerationalId<StringId>> {
88        let id = self.resolve(s)?;
89        Ok(GenerationalId::new(id, self.generation))
90    }
91
92    pub fn intern_managed_strings(
93        &mut self,
94        s: &[ManagedStringId],
95        out: &mut [GenerationalId<StringId>],
96    ) -> anyhow::Result<()> {
97        anyhow::ensure!(s.len() == out.len());
98        for i in 0..s.len() {
99            out[i] = self.intern_managed_string(s[i])?;
100        }
101        Ok(())
102    }
103
104    pub fn intern_mapping(
105        &mut self,
106        memory_start: u64,
107        memory_limit: u64,
108        file_offset: u64,
109        filename: GenerationalId<StringId>,
110        build_id: GenerationalId<StringId>,
111    ) -> anyhow::Result<GenerationalId<MappingId>> {
112        let mapping = Mapping {
113            memory_start,
114            memory_limit,
115            file_offset,
116            filename: filename.get(self.generation)?,
117            build_id: build_id.get(self.generation)?,
118        };
119        let id = self.mappings.dedup(mapping);
120        Ok(GenerationalId::new(id, self.generation))
121    }
122
123    pub fn intern_sample(
124        &mut self,
125        stacktrace: GenerationalId<StackTraceId>,
126        values: &[i64],
127        labels: GenerationalId<LabelSetId>,
128        timestamp: Option<Timestamp>,
129    ) -> anyhow::Result<()> {
130        // TODO: validate sample labels? Or should we do that when we make the label set?
131        anyhow::ensure!(
132            values.len() == self.sample_types.len(),
133            "expected {} sample types, but sample had {} sample types",
134            self.sample_types.len(),
135            values.len(),
136        );
137        let stacktrace = stacktrace.get(self.generation)?;
138        let labels = labels.get(self.generation)?;
139
140        self.observations
141            .add(Sample::new(labels, stacktrace), timestamp, values)
142    }
143
144    pub fn intern_stacktrace(
145        &mut self,
146        locations: &[GenerationalId<LocationId>],
147    ) -> anyhow::Result<GenerationalId<StackTraceId>> {
148        let locations = locations
149            .iter()
150            .map(|l| l.get(self.generation))
151            .collect::<anyhow::Result<Vec<_>>>()?;
152        let stacktrace = StackTrace { locations };
153        let id = self.stack_traces.dedup(stacktrace);
154        Ok(GenerationalId::new(id, self.generation))
155    }
156
157    pub const INTERNED_EMPTY_STRING: GenerationalId<StringId> =
158        GenerationalId::new_immortal(StringId::ZERO);
159
160    pub fn intern_string(&mut self, s: &str) -> anyhow::Result<GenerationalId<StringId>> {
161        if s.is_empty() {
162            Ok(Self::INTERNED_EMPTY_STRING)
163        } else {
164            Ok(GenerationalId::new(self.try_intern(s)?, self.generation))
165        }
166    }
167
168    pub fn intern_strings(
169        &mut self,
170        s: &[&str],
171        out: &mut [GenerationalId<StringId>],
172    ) -> anyhow::Result<()> {
173        anyhow::ensure!(s.len() == out.len());
174        for i in 0..s.len() {
175            out[i] = self.intern_string(s[i])?;
176        }
177        Ok(())
178    }
179
180    // Simple synchronization between samples and profile rotation/export.
181    // Interning a sample may require several calls to the profiler to intern intermediate values,
182    // which are not inherently atomic.  Since these intermediate values are tied to a particular
183    // profiler generation, and are invalidated when the generation changes, some coordination must
184    // occur between sampling and profile rotation/export.
185    // When the generation changes, one of three things can happen:
186    // 1. The sample can be dropped.
187    // 2. The sample can be recreated and interned into the new profile.
188    // 3. The profile rotation should wait until the sampling operation is complete.
189    //
190    // This API provides a mechanism for samples to pause rotation until they complete, and
191    // for samples to be notified that a rotation is in progress so they can wait to begin.
192    // There are probably better ways, and maybe we should have a notification mechanism.
193    // But for now this should be enough.
194    const FLAG: u64 = u32::MAX as u64;
195
196    /// Prevent any new samples from starting.
197    /// Returns the number of remaining samples.
198    pub fn sample_block(&mut self) -> anyhow::Result<u64> {
199        let current = self.active_samples.fetch_add(Self::FLAG, SeqCst);
200        if current >= Self::FLAG {
201            self.active_samples.fetch_sub(Self::FLAG, SeqCst);
202        }
203        Ok(current % Self::FLAG)
204    }
205
206    pub fn sample_end(&mut self) -> anyhow::Result<()> {
207        self.active_samples.fetch_sub(1, SeqCst);
208        Ok(())
209    }
210
211    pub fn sample_start(&mut self) -> anyhow::Result<()> {
212        let old = self.active_samples.fetch_add(1, SeqCst);
213        if old >= Self::FLAG {
214            self.active_samples.fetch_sub(1, SeqCst);
215            anyhow::bail!("Can't start sample, export in progress");
216        }
217        Ok(())
218    }
219
220    pub fn samples_active(&mut self) -> anyhow::Result<u64> {
221        let current = self.active_samples.load(SeqCst);
222        Ok(current % Self::FLAG)
223    }
224
225    pub fn samples_are_blocked(&mut self) -> anyhow::Result<bool> {
226        let current = self.active_samples.load(SeqCst);
227        Ok(current >= Self::FLAG)
228    }
229
230    pub fn samples_are_drained(&mut self) -> anyhow::Result<bool> {
231        let current = self.active_samples.load(SeqCst);
232        Ok(current % Self::FLAG == 0)
233    }
234}