From 041231e7bb2af2b185251515440bcf96042bdf2c Mon Sep 17 00:00:00 2001 From: Levi Morrison Date: Tue, 15 Aug 2023 04:19:31 -0600 Subject: [PATCH] extract small_non_zero_pprof_id --- profiling/src/profile/internal/function.rs | 13 +++---------- profiling/src/profile/internal/location.rs | 14 +++----------- profiling/src/profile/internal/mapping.rs | 13 +++---------- profiling/src/profile/internal/mod.rs | 9 +++++++++ 4 files changed, 18 insertions(+), 31 deletions(-) diff --git a/profiling/src/profile/internal/function.rs b/profiling/src/profile/internal/function.rs index 12bdd1687..c27303269 100644 --- a/profiling/src/profile/internal/function.rs +++ b/profiling/src/profile/internal/function.rs @@ -2,7 +2,7 @@ // This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2021-Present Datadog, Inc. use super::super::{pprof, StringId}; -use super::{Id, Item, PprofItem}; +use super::{small_non_zero_pprof_id, Id, Item, PprofItem}; use std::fmt::Debug; use std::num::NonZeroU32; @@ -39,15 +39,8 @@ pub struct FunctionId(NonZeroU32); impl Id for FunctionId { type RawId = u64; - fn from_offset(v: usize) -> Self { - let index: u32 = v.try_into().expect("FunctionId to fit into a u32"); - - // PProf reserves function 0. - // Both this, and the serialization of the table, add 1 to avoid the 0 element - let index = index.checked_add(1).expect("FunctionId to fit into a u32"); - // Safety: the `checked_add(1).expect(...)` guards this from ever being zero. - let index = unsafe { NonZeroU32::new_unchecked(index) }; - Self(index) + fn from_offset(offset: usize) -> Self { + Self(small_non_zero_pprof_id(offset).expect("FunctionId to fit into a u32")) } fn to_raw_id(&self) -> Self::RawId { diff --git a/profiling/src/profile/internal/location.rs b/profiling/src/profile/internal/location.rs index aa889e823..acf2cb7d9 100644 --- a/profiling/src/profile/internal/location.rs +++ b/profiling/src/profile/internal/location.rs @@ -2,8 +2,7 @@ // This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2021-Present Datadog, Inc. use super::super::{pprof, MappingId}; -use super::{Id, Item, Line, PprofItem}; -use std::fmt::Debug; +use super::{small_non_zero_pprof_id, Id, Item, Line, PprofItem}; use std::num::NonZeroU32; #[derive(Clone, Debug, Eq, PartialEq, Hash, PartialOrd, Ord)] @@ -39,15 +38,8 @@ pub struct LocationId(NonZeroU32); impl Id for LocationId { type RawId = u64; - fn from_offset(v: usize) -> Self { - let index: u32 = v.try_into().expect("LocationId to fit into a u32"); - - // PProf reserves location 0. - // Both this, and the serialization of the table, add 1 to avoid the 0 element - let index = index.checked_add(1).expect("LocationId to fit into a u32"); - // Safety: the `checked_add(1).expect(...)` guards this from ever being zero. - let index = unsafe { NonZeroU32::new_unchecked(index) }; - Self(index) + fn from_offset(offset: usize) -> Self { + Self(small_non_zero_pprof_id(offset).expect("LocationId to fit into a u32")) } fn to_raw_id(&self) -> Self::RawId { diff --git a/profiling/src/profile/internal/mapping.rs b/profiling/src/profile/internal/mapping.rs index c78a168e0..67aba9940 100644 --- a/profiling/src/profile/internal/mapping.rs +++ b/profiling/src/profile/internal/mapping.rs @@ -1,7 +1,7 @@ // Unless explicitly stated otherwise all files in this repository are licensed under the Apache License Version 2.0. // This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2023-Present Datadog, Inc. -use super::{Id, Item, PprofItem}; +use super::{small_non_zero_pprof_id, Id, Item, PprofItem}; use crate::profile::{pprof, StringId}; use std::num::NonZeroU32; @@ -52,15 +52,8 @@ pub struct MappingId(NonZeroU32); impl Id for MappingId { type RawId = u64; - fn from_offset(v: usize) -> Self { - let index: u32 = v.try_into().expect("MappingId to fit into a u32"); - - // PProf reserves location 0. - // Both this, and the serialization of the table, add 1 to avoid the 0 element - let index = index.checked_add(1).expect("MappingId to fit into a u32"); - // Safety: the `checked_add(1).expect(...)` guards this from ever being zero. - let index = unsafe { NonZeroU32::new_unchecked(index) }; - Self(index) + fn from_offset(offset: usize) -> Self { + Self(small_non_zero_pprof_id(offset).expect("MappingId to fit into a u32")) } fn to_raw_id(&self) -> Self::RawId { diff --git a/profiling/src/profile/internal/mod.rs b/profiling/src/profile/internal/mod.rs index 2937c87b6..3507ec655 100644 --- a/profiling/src/profile/internal/mod.rs +++ b/profiling/src/profile/internal/mod.rs @@ -20,6 +20,7 @@ pub use string::*; pub use value_type::*; use std::hash::Hash; +use std::num::NonZeroU32; pub trait Id: Copy + Eq + Hash { type RawId; @@ -54,3 +55,11 @@ pub trait PprofItem: Item { fn to_pprof(&self, id: Self::Id) -> Self::PprofMessage; } + +#[inline] +fn small_non_zero_pprof_id(offset: usize) -> Option { + let small: u32 = offset.try_into().ok()?; + let non_zero = small.checked_add(1)?; + // Safety: the `checked_add(1)?` guards this from ever being zero. + Some(unsafe { NonZeroU32::new_unchecked(non_zero) }) +}