[PROF-11405] Fix managed string storage causing corrupt profiles

profiling/src/collections/string_storage.rs

@@ -1,11 +1,11 @@
 // Copyright 2023-Present Datadog, Inc. https://www.datadoghq.com/
 // SPDX-License-Identifier: Apache-2.0
 
+use anyhow::Context;
 use std::cell::Cell;
 use std::collections::HashMap;
 use std::hash::BuildHasherDefault;
 use std::num::NonZeroU32;
-use std::ptr;
 use std::rc::Rc;
 
 use super::identifiable::StringId;
@@ -14,7 +14,7 @@ use super::string_table::StringTable;
 #[derive(PartialEq, Debug)]
 struct ManagedStringData {
     str: Rc<str>,
-    cached_seq_num: Cell<Option<(*const StringTable, StringId)>>,
+    cached_seq_num: Cell<Option<(InternalCachedProfileId, StringId)>>,
     usage_count: Cell<u32>,
 }
 
@@ -23,6 +23,40 @@ pub struct ManagedStringStorage {
     id_to_data: HashMap<u32, ManagedStringData, BuildHasherDefault<rustc_hash::FxHasher>>,
     str_to_id: HashMap<Rc<str>, u32, BuildHasherDefault<rustc_hash::FxHasher>>,
     current_gen: u32,
+    next_cached_profile_id: InternalCachedProfileId,
+}
+
+#[derive(PartialEq, Debug)]
+// The `ManagedStringStorage::get_seq_num` operation is used to map a `ManagedStorageId` into a
+// `StringId` for a given `StringTable`. As an optimization, we store a one-element `cached_seq_num`
+//  inline cache with each `ManagedStringData` entry, so that repeatedly calling
+// `get_seq_num` with the same id provides faster lookups.
+//
+// Because:
+// 1. Multiple profiles can be using the same `ManagedStringTable`
+// 2. The same profile resets its `StringTable` on serialization and starts anew
+// ...we need a way to identify when the cache can and cannot be reused.
+//
+// This is where the `CachedProfileId` comes in. A given `CachedProfileId` should be considered
+// as representing a unique `StringTable`. Different `StringTable`s should have different
+// `CachedProfileId`s, and when a `StringTable` gets flushed and starts anew, it should also have a
+// different `CachedProfileId`.
+//
+// **This struct is on purpose not Copy and not Clone to try to make it really hard to accidentally
+// reuse** when a profile gets reset.
+pub struct CachedProfileId {
+    id: u32,
+}
+
+#[derive(PartialEq, Debug, Copy, Clone)]
+struct InternalCachedProfileId {
+    id: u32,
+}
+
+impl From<&CachedProfileId> for InternalCachedProfileId {
+    fn from(cached: &CachedProfileId) -> Self {
+        InternalCachedProfileId { id: cached.id }
+    }
 }
 
 impl ManagedStringStorage {
@@ -32,13 +66,24 @@ impl ManagedStringStorage {
             id_to_data: Default::default(),
             str_to_id: Default::default(),
             current_gen: 0,
+            next_cached_profile_id: InternalCachedProfileId { id: 0 },
         };
         // Ensure empty string gets id 0 and always has usage > 0 so it's always retained
         // Safety: On an empty managed string table intern should never fail.
         storage.intern_new("").expect("Initialization to succeed");
         storage
     }
 
+    pub fn next_cached_profile_id(&mut self) -> anyhow::Result<CachedProfileId> {
+        let next_id = self.next_cached_profile_id.id;
+        self.next_cached_profile_id = InternalCachedProfileId {
+            id: next_id
+                .checked_add(1)
+                .context("Ran out of cached_profile_ids!")?,
+        };
+        Ok(CachedProfileId { id: next_id })
+    }
+
     pub fn advance_gen(&mut self) {
         self.id_to_data.retain(|_, data| {
             let retain = data.usage_count.get() > 0;
@@ -103,21 +148,21 @@ impl ManagedStringStorage {
 
     // Here id is a NonZeroU32 because an id of 0 which StringTable always maps to 0 as well so this
     // entire call can be skipped
+    // See comment on `struct CachedProfileId` for details on how to use it.
     pub fn get_seq_num(
         &self,
         id: NonZeroU32,
         profile_strings: &mut StringTable,
+        cached_profile_id: &CachedProfileId,
     ) -> anyhow::Result<StringId> {
         let data = self.get_data(id.into())?;
 
-        let profile_strings_pointer = ptr::addr_of!(*profile_strings);
-
         match data.cached_seq_num.get() {
-            Some((pointer, seq_num)) if pointer == profile_strings_pointer => Ok(seq_num),
+            Some((profile_id, seq_num)) if profile_id.id == cached_profile_id.id => Ok(seq_num),
             _ => {
                 let seq_num = profile_strings.intern(data.str.as_ref());
                 data.cached_seq_num
-                    .set(Some((profile_strings_pointer, seq_num)));
+                    .set(Some((cached_profile_id.into(), seq_num)));
                 Ok(seq_num)
             }
         }

profiling/src/internal/profile/mod.rs

@@ -9,7 +9,7 @@ use super::*;
 use crate::api;
 use crate::api::ManagedStringId;
 use crate::collections::identifiable::*;
-use crate::collections::string_storage::ManagedStringStorage;
+use crate::collections::string_storage::{CachedProfileId, ManagedStringStorage};
 use crate::collections::string_table::StringTable;
 use crate::iter::{IntoLendingIterator, LendingIterator};
 use crate::pprof::sliced_proto::*;
@@ -43,6 +43,7 @@ pub struct Profile {
     start_time: SystemTime,
     strings: StringTable,
     string_storage: Option<Rc<RwLock<ManagedStringStorage>>>,
+    string_storage_cached_profile_id: Option<CachedProfileId>,
     timestamp_key: StringId,
     upscaling_rules: UpscalingRules,
 }
@@ -199,14 +200,33 @@ impl Profile {
             return Ok(StringId::ZERO); // Both string tables use zero for the empty string
         };
 
-        self.string_storage
+        let string_storage = self.string_storage
             .as_ref()
             // Safety: We always get here through a direct or indirect call to add_string_id_sample,
             // which already ensured that the string storage exists.
-            .ok_or_else(|| anyhow::anyhow!("Current sample makes use of ManagedStringIds but profile was not created using a string table"))?
+            .ok_or_else(|| anyhow::anyhow!("Current sample makes use of ManagedStringIds but profile was not created using a string table"))?;
+
+        let cached_profile_id = match self.string_storage_cached_profile_id.as_ref() {
+            Some(cached_profile_id) => cached_profile_id,
+            None => {
+                let new_id = string_storage
+                    .write()
+                    .map_err(|_| {
+                        anyhow::anyhow!(
+                            "acquisition of write lock on string storage should succeed"
+                        )
+                    })?
+                    .next_cached_profile_id()?;
+                self.string_storage_cached_profile_id.get_or_insert(new_id)
+            }
+        };
+
+        string_storage
             .read()
-            .map_err(|_| anyhow::anyhow!("acquisition of read lock on string storage should succeed"))?
-            .get_seq_num(non_empty_string_id, &mut self.strings)
+            .map_err(|_| {
+                anyhow::anyhow!("acquisition of read lock on string storage should succeed")
+            })?
+            .get_seq_num(non_empty_string_id, &mut self.strings, cached_profile_id)
     }
 
     /// Creates a profile with `start_time`.
@@ -577,6 +597,8 @@ impl Profile {
             start_time,
             strings: Default::default(),
             string_storage,
+            string_storage_cached_profile_id: None, /* Never reuse an id! See comments on
+                                                     * CachedProfileId for why. */
             timestamp_key: Default::default(),
             upscaling_rules: Default::default(),
         };
@@ -2277,4 +2299,67 @@ mod api_tests {
         }
         Ok(())
     }
+
+    #[test]
+    fn test_regression_managed_string_table_correctly_maps_ids() {
+        let storage = Rc::new(RwLock::new(ManagedStringStorage::new()));
+        let hello_id: u32;
+        let world_id: u32;
+
+        {
+            let mut storage_guard = storage.write().unwrap();
+            hello_id = storage_guard.intern("hello").unwrap();
+            world_id = storage_guard.intern("world").unwrap();
+        }
+
+        let sample_types = [api::ValueType::new("samples", "count")];
+        let mut profile =
+            Profile::with_string_storage(SystemTime::now(), &sample_types, None, storage.clone());
+
+        let location = api::StringIdLocation {
+            function: api::StringIdFunction {
+                name: api::ManagedStringId { value: hello_id },
+                filename: api::ManagedStringId { value: world_id },
+                ..Default::default()
+            },
+            ..Default::default()
+        };
+
+        let sample = api::StringIdSample {
+            locations: vec![location],
+            values: vec![1],
+            labels: vec![],
+        };
+
+        profile.add_string_id_sample(sample.clone(), None).unwrap();
+        profile.add_string_id_sample(sample.clone(), None).unwrap();
+
+        let pprof_first_profile =
+            pprof::roundtrip_to_pprof(profile.reset_and_return_previous(None).unwrap()).unwrap();
+
+        assert!(pprof_first_profile
+            .string_table
+            .iter()
+            .any(|s| s == "hello"));
+        assert!(pprof_first_profile
+            .string_table
+            .iter()
+            .any(|s| s == "world"));
+
+        // If the cache invalidation on the managed string table is working correctly, these strings
+        // get correctly re-added to the profile's string table
+
+        profile.add_string_id_sample(sample.clone(), None).unwrap();
+        profile.add_string_id_sample(sample.clone(), None).unwrap();
+        let pprof_second_profile = pprof::roundtrip_to_pprof(profile).unwrap();
+
+        assert!(pprof_second_profile
+            .string_table
+            .iter()
+            .any(|s| s == "hello"));
+        assert!(pprof_second_profile
+            .string_table
+            .iter()
+            .any(|s| s == "world"));
+    }
 }