Skip to main content

micromegas_tracing/
dyn_deps.rs

1//! Per-block dependency records for runtime property values, kept in their wire form.
2use crate::{
3    property_set::{
4        DynProperty, Property, PropertyDepsQueue, PropertySet, PropertySetDependency,
5        PropertySetId, StringDepsQueue, cmp_dyn_properties,
6    },
7    static_string_ref::{BlockStringRef, StaticStringRef},
8};
9use micromegas_transit::{
10    HeterogeneousQueue, StaticStringDependency, read_any, string_codec::StringCodec,
11};
12use rustc_hash::{FxHashMap, FxHashSet};
13use std::{borrow::Cow, collections::hash_map::Entry, fmt};
14use xxhash_rust::xxh3::{xxh3_64, xxh3_64_with_seed};
15
16/// Keeps dynamic ids out of the range of pointer ids, which the parser requires to be unique
17/// within a block. No supported platform hands out addresses with this bit set.
18pub const DYN_ID_TAG: u64 = 1 << 63;
19
20const FIRST_INSERT_CAPACITY: usize = 64 * 1024;
21
22// Each record is a type byte and a u32 size, followed by the payload.
23const RECORD_HEADER: usize = 1 + std::mem::size_of::<u32>();
24// String payload: id, codec, len, then the bytes.
25const STRING_HEADER: usize = std::mem::size_of::<u64>() + 1 + std::mem::size_of::<u32>();
26// Set payload: id, count, then the properties.
27const SET_HEADER: usize = std::mem::size_of::<u64>() + std::mem::size_of::<u32>();
28
29fn dyn_string_id(offset: u32) -> u64 {
30    DYN_ID_TAG | (u64::from(offset) << 1)
31}
32
33fn dyn_set_id(offset: u32) -> u64 {
34    DYN_ID_TAG | (u64::from(offset) << 1) | 1
35}
36
37fn dyn_offset(id: u64) -> u32 {
38    ((id & !DYN_ID_TAG) >> 1) as u32
39}
40
41/// The result of combining a callsite's static set with runtime properties.
42pub struct MergedProperties<'a> {
43    pub properties: Vec<DynProperty<'a>>,
44    // Bit i set: properties[i].value comes from the `&'static` set.
45    static_values: u64,
46    // Bit i set: properties[i].name is `'static`.
47    static_names: u64,
48}
49
50impl MergedProperties<'_> {
51    fn is_static_value(&self, index: usize) -> bool {
52        is_static_bit(self.static_values, index)
53    }
54
55    fn is_static_name(&self, index: usize) -> bool {
56        is_static_bit(self.static_names, index)
57    }
58}
59
60fn is_static_bit(mask: u64, index: usize) -> bool {
61    index < 64 && (mask >> index) & 1 != 0
62}
63
64/// Combines a static set with dynamic properties into one sorted list.
65///
66/// A dynamic value replaces a static property of the same name (compared case-insensitively
67/// to match `property_get`).
68#[inline]
69pub fn merge_properties<'a>(
70    static_set: &'static PropertySet,
71    dyn_properties: &[DynProperty<'a>],
72) -> MergedProperties<'a> {
73    let mut merged: Vec<DynProperty<'a>> = static_set.dyn_properties().to_vec();
74    for dyn_prop in dyn_properties {
75        merged.retain(|p| !p.name().eq_ignore_ascii_case(dyn_prop.name()));
76        merged.push(*dyn_prop);
77    }
78    merged.sort_by(cmp_dyn_properties);
79    // A value that is the very memory of a static property's value is `'static` whoever
80    // passed it. This avoids carrying a flag through the retain and the sort.
81    let (static_values, static_names) =
82        merged
83            .iter()
84            .take(64)
85            .enumerate()
86            .fold((0u64, 0u64), |(values, names), (i, p)| {
87                let is_static = static_set.get_properties().iter().any(|s| {
88                    s.value.id() == p.value().as_ptr() as u64
89                        && s.value.len as usize == p.value().len()
90                        && s.name.as_str() == p.name()
91                });
92                (
93                    values | (u64::from(is_static) << i),
94                    names | (u64::from(p.is_static_name()) << i),
95                )
96            });
97    MergedProperties {
98        properties: merged,
99        static_values,
100        static_names,
101    }
102}
103
104fn content_hash<'a>(properties: impl Iterator<Item = (&'a str, &'a str)>) -> u64 {
105    properties.fold(0, |h, (name, value)| {
106        xxh3_64_with_seed(value.as_bytes(), xxh3_64_with_seed(name.as_bytes(), h))
107    })
108}
109
110#[inline]
111fn read_u32(bytes: &[u8], at: usize) -> u32 {
112    u32::from_ne_bytes(bytes[at..at + 4].try_into().unwrap())
113}
114
115#[inline]
116fn string_at(bytes: &[u8], offset: u32) -> &str {
117    let payload = offset as usize + RECORD_HEADER;
118    let len = read_u32(bytes, payload + std::mem::size_of::<u64>() + 1) as usize;
119    let begin = payload + STRING_HEADER;
120    // SAFETY: the bytes were copied from a `&str` by `intern_str`.
121    unsafe { std::str::from_utf8_unchecked(&bytes[begin..begin + len]) }
122}
123
124#[inline]
125fn set_len(bytes: &[u8], offset: u32) -> usize {
126    read_u32(
127        bytes,
128        offset as usize + RECORD_HEADER + std::mem::size_of::<u64>(),
129    ) as usize
130}
131
132#[inline]
133fn set_property(bytes: &[u8], offset: u32, index: usize) -> Property {
134    let size = std::mem::size_of::<Property>();
135    let begin = offset as usize + RECORD_HEADER + SET_HEADER + index * size;
136    let window = &bytes[begin..begin + size];
137    // SAFETY: the window holds a `Property` written by `write_any`.
138    unsafe { read_any::<Property>(window.as_ptr()) }
139}
140
141/// Resolves a ref that was built from a `&'static str`.
142///
143/// # Safety
144/// `r` must have been built from a `&'static str`.
145pub(crate) unsafe fn static_ref_str(r: &StaticStringRef) -> &'static str {
146    unsafe {
147        let slice = std::slice::from_raw_parts(r.ptr as *const u8, r.len as usize);
148        std::str::from_utf8_unchecked(slice)
149    }
150}
151
152/// Records keyed by content hash. A collision is vanishingly rare, so the common case stores
153/// one record inline and allocates nothing per key.
154struct HashIndex<T> {
155    first: FxHashMap<u64, T>,
156    collisions: FxHashMap<u64, Vec<T>>,
157}
158
159impl<T> HashIndex<T> {
160    fn new() -> Self {
161        Self {
162            first: FxHashMap::default(),
163            collisions: FxHashMap::default(),
164        }
165    }
166
167    fn candidates(&self, hash: u64) -> impl Iterator<Item = &T> {
168        let extra = if self.collisions.is_empty() {
169            None
170        } else {
171            self.collisions.get(&hash)
172        };
173        self.first
174            .get(&hash)
175            .into_iter()
176            .chain(extra.into_iter().flatten())
177    }
178
179    fn insert(&mut self, hash: u64, record: T) {
180        match self.first.entry(hash) {
181            Entry::Vacant(slot) => {
182                slot.insert(record);
183            }
184            Entry::Occupied(_) => self.collisions.entry(hash).or_default().push(record),
185        }
186    }
187
188    fn iter(&self) -> impl Iterator<Item = &T> {
189        self.first
190            .values()
191            .chain(self.collisions.values().flatten())
192    }
193}
194
195#[derive(Debug)]
196struct DynSet {
197    offset: u32,
198    // The masks of the `MergedProperties` the record was built with.
199    static_values: u64,
200    static_names: u64,
201}
202
203/// The dependency records of one block that only exist at run time.
204///
205/// Records are serialized when the event is pushed, with ids that are scoped to the block, and
206/// appended as-is after the static dependencies when the block is extracted. A record only
207/// references records pushed before it, or `'static` data.
208pub struct DynDeps<D> {
209    queue: D,
210    // offsets of string records
211    strings: HashIndex<u32>,
212    sets: HashIndex<DynSet>,
213}
214
215impl<D> fmt::Debug for DynDeps<D> {
216    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
217        f.debug_struct("DynDeps").finish_non_exhaustive()
218    }
219}
220
221impl<D: HeterogeneousQueue> Default for DynDeps<D> {
222    fn default() -> Self {
223        Self::new()
224    }
225}
226
227impl<D: HeterogeneousQueue> DynDeps<D> {
228    /// Allocates nothing until the first insert.
229    pub fn new() -> Self {
230        Self {
231            queue: D::new(0),
232            strings: HashIndex::new(),
233            sets: HashIndex::new(),
234        }
235    }
236
237    /// Exact size of the dependency bytes this adds to the encoded block.
238    pub fn len_bytes(&self) -> usize {
239        self.queue.len_bytes()
240    }
241
242    pub fn queue(&self) -> &D {
243        &self.queue
244    }
245
246    fn reserve_on_first_insert(&mut self) {
247        if self.queue.capacity_bytes() == 0 {
248            self.queue = D::new(FIRST_INSERT_CAPACITY);
249        }
250    }
251
252    /// Decodes a set for in-process readers.
253    pub fn set_properties(&self, id: PropertySetId) -> Vec<(&str, &str)> {
254        if let Some(set) = id.as_interned() {
255            return set
256                .get_properties()
257                .iter()
258                .map(|p| (p.name.as_str(), p.value.as_str()))
259                .collect();
260        }
261        let offset = dyn_offset(id.0);
262        let bytes = self.queue.as_bytes();
263        let set = self
264            .sets
265            .iter()
266            .find(|s| s.offset == offset)
267            .expect("unknown dynamic set id");
268        (0..set_len(bytes, offset))
269            .map(|i| {
270                let prop = set_property(bytes, offset, i);
271                (
272                    ref_str(bytes, &prop.name, set.static_names, i),
273                    ref_str(bytes, &prop.value, set.static_values, i),
274                )
275            })
276            .collect()
277    }
278
279    /// The text of a string record written by `intern_str`.
280    pub fn str(&self, id: u64) -> &str {
281        string_at(self.queue.as_bytes(), dyn_offset(id))
282    }
283}
284
285/// Resolves a name or value of a dynamic set: a `'static` string when its bit is set in
286/// `mask`, a string record of the block otherwise.
287#[inline]
288fn ref_str<'b>(bytes: &'b [u8], r: &StaticStringRef, mask: u64, index: usize) -> &'b str {
289    if is_static_bit(mask, index) {
290        // SAFETY: masked refs were built from `&'static str`s.
291        unsafe { static_ref_str(r) }
292    } else {
293        string_at(bytes, dyn_offset(r.ptr))
294    }
295}
296
297impl<D: StringDepsQueue> DynDeps<D> {
298    /// Returns the id of the string record holding `value`, pushing one on first sight.
299    pub fn intern_str(&mut self, value: &str) -> u64 {
300        self.reserve_on_first_insert();
301        let hash = xxh3_64(value.as_bytes());
302        let bytes = self.queue.as_bytes();
303        for &offset in self.strings.candidates(hash) {
304            if string_at(bytes, offset) == value {
305                return dyn_string_id(offset);
306            }
307        }
308        let offset = u32::try_from(self.queue.len_bytes()).expect("dynamic queue exceeds 4 GiB");
309        let id = dyn_string_id(offset);
310        self.queue.push_string(StaticStringDependency {
311            id,
312            codec: StringCodec::Utf8,
313            len: value.len() as u32,
314            ptr: value.as_ptr(),
315        });
316        self.strings.insert(hash, offset);
317        id
318    }
319
320    /// Like `intern_str`, as the reference an event stores.
321    pub fn intern_str_ref(&mut self, value: &str) -> BlockStringRef {
322        BlockStringRef {
323            id: self.intern_str(value),
324            len: value.len() as u32,
325            codec: StringCodec::Utf8,
326        }
327    }
328}
329
330impl<D: PropertyDepsQueue> DynDeps<D> {
331    /// Returns the id of the set record equal to `merged`, pushing it (and the string records
332    /// it needs) on first sight. A hit allocates nothing.
333    #[inline]
334    pub fn intern_set(&mut self, merged: &MergedProperties<'_>) -> PropertySetId {
335        self.reserve_on_first_insert();
336        let hash = content_hash(merged.properties.iter().map(|p| (p.name(), p.value())));
337        let bytes = self.queue.as_bytes();
338        for candidate in self.sets.candidates(hash) {
339            if Self::same_content(bytes, candidate, merged) {
340                return PropertySetId(dyn_set_id(candidate.offset));
341            }
342        }
343        let mut properties = Vec::with_capacity(merged.properties.len());
344        for (i, prop) in merged.properties.iter().enumerate() {
345            // Only ever serialized: the ptr of a dynamic ref is the id of the string record.
346            let name = if merged.is_static_name(i) {
347                // SAFETY: the mask marks names that were built from a `&'static str`.
348                unsafe { StaticStringRef::from_str_unchecked(prop.name()) }
349            } else {
350                StaticStringRef {
351                    ptr: self.intern_str(prop.name()),
352                    len: prop.name().len() as u32,
353                    codec: StringCodec::Utf8,
354                }
355            };
356            let value = if merged.is_static_value(i) {
357                // SAFETY: the mask marks values that came from a `&'static` set.
358                unsafe { StaticStringRef::from_str_unchecked(prop.value()) }
359            } else {
360                StaticStringRef {
361                    ptr: self.intern_str(prop.value()),
362                    len: prop.value().len() as u32,
363                    codec: StringCodec::Utf8,
364                }
365            };
366            properties.push(Property { name, value });
367        }
368        let offset = u32::try_from(self.queue.len_bytes()).expect("dynamic queue exceeds 4 GiB");
369        let id = dyn_set_id(offset);
370        self.queue.push_set(PropertySetDependency {
371            id,
372            properties: Cow::Owned(properties),
373        });
374        self.sets.insert(
375            hash,
376            DynSet {
377                offset,
378                static_values: merged.static_values,
379                static_names: merged.static_names,
380            },
381        );
382        PropertySetId(id)
383    }
384
385    #[inline]
386    fn same_content(bytes: &[u8], candidate: &DynSet, merged: &MergedProperties<'_>) -> bool {
387        if set_len(bytes, candidate.offset) != merged.properties.len() {
388            return false;
389        }
390        merged.properties.iter().enumerate().all(|(i, wanted)| {
391            let prop = set_property(bytes, candidate.offset, i);
392            ref_str(bytes, &prop.name, candidate.static_names, i) == wanted.name()
393                && ref_str(bytes, &prop.value, candidate.static_values, i) == wanted.value()
394        })
395    }
396
397    /// Emits the `'static` names and values that dynamic sets reference, once each. Block-scoped
398    /// ones are already in the queue.
399    pub(crate) fn record_static_refs(&self, recorded_deps: &mut FxHashSet<u64>, deps: &mut D) {
400        let bytes = self.queue.as_bytes();
401        for set in self.sets.iter() {
402            for i in 0..set_len(bytes, set.offset) {
403                let prop = set_property(bytes, set.offset, i);
404                if is_static_bit(set.static_names, i) && recorded_deps.insert(prop.name.id()) {
405                    deps.push_string(prop.name.into_dependency());
406                }
407                if is_static_bit(set.static_values, i) && recorded_deps.insert(prop.value.id()) {
408                    deps.push_string(prop.value.into_dependency());
409                }
410            }
411        }
412    }
413}