Add store_data + page_out_item to PagableStorage trait

Summary:
Adds two methods to the `PagableStorage` trait so that callers can use
`Arc<dyn PagableStorage>` rather than the concrete sled type:

- `store_data(key, data)`: backend-specific primitive — writes one
  content-addressable blob.
- `page_out_item(serializer, finished)`: default trait method that walks
  the arc tree, recursively serializes nested arcs, and stores them via
  `store_data`.

The orchestration logic (previously inherent on `SledBackedPagableStorage`)
is moved into the default trait method, so adding a new backend only requires
implementing `store_data`. `SledBackedPagableStorage`, `InMemoryPagableStorageHandle`,
and the test `EmptyPagableStorage` each gain a `store_data` impl.

`SerializerForPaging::new` is also bumped from `pub(crate)` to `pub` so the
default trait method is callable from outside the crate.

No behavior change beyond what's described above.

Reviewed By: jtbraun

Differential Revision: D101759762

fbshipit-source-id: 6193bb529dd17d1030cb3eb444d72413a73bbb48

Author: cjhopman

pagable/src/storage/in_memory.rs

@@ -294,4 +294,10 @@ impl PagableStorage for InMemoryPagableStorageHandle {
     fn session_context(&self) -> &Mutex<SessionContext> {
         &self.session_context
     }
+
+    fn store_data(&self, data: PagableData) -> DataKey {
+        let key = data.compute_key();
+        self.cache.insert_data(key, std::sync::Arc::new(data));
+        key
+    }
 }

pagable/src/storage/sled.rs

@@ -75,33 +75,6 @@ impl SledBackedPagableStorage {
         lock.pending.len()
     }
 
-    /// Serialize `PagableData` into the on-disk byte format and insert into sled.
-    /// Returns the content-addressable `DataKey`.
-    fn store_data(&self, data: PagableData) -> DataKey {
-        let key = data.compute_key();
-        let db_key = bytemuck::bytes_of(&key);
-        if self
-            .db
-            .contains_key(db_key)
-            .expect("sled contains_key failed")
-        {
-            return key;
-        }
-
-        let bytes_size = 8 + 8 + data.data.len() + data.arcs.len() * 16;
-        let mut bytes = Vec::with_capacity(bytes_size);
-        bytes.extend_from_slice(&(data.data.len() as u64).to_le_bytes());
-        bytes.extend_from_slice(&(data.arcs.len() as u64).to_le_bytes());
-        bytes.extend_from_slice(&data.data);
-        bytes.extend_from_slice(bytemuck::cast_slice(&data.arcs));
-        assert_eq!(bytes.len(), bytes_size);
-
-        self.db
-            .insert(db_key, sled::IVec::from(bytes))
-            .expect("sled insert failed");
-        key
-    }
-
     /// Recursively serialize a set of arcs and their dependencies, storing each
     /// into sled. Updates `finished` with the identity -> DataKey mapping for
     /// each processed arc.
@@ -194,42 +167,6 @@ impl SledBackedPagableStorage {
         }
     }
 
-    pub fn page_out_item_cache(&self) -> HashMap<usize, DataKey> {
-        HashMap::new()
-    }
-
-    pub fn serializer_for_page_out_item<'a>(
-        &self,
-        session_context: &'a mut SessionContext,
-    ) -> SerializerForPaging<'a> {
-        SerializerForPaging::new(session_context)
-    }
-
-    pub fn page_out_item(
-        &self,
-        item_data: Vec<u8>,
-        item_arcs: Vec<Box<dyn ArcEraseDyn>>,
-        finished: &mut HashMap<usize, DataKey>,
-        session_context: &mut SessionContext,
-    ) -> DataKey {
-        let roots = item_arcs.iter().map(|arc| arc.clone_dyn()).collect();
-        self.serialize_arcs(roots, finished, session_context);
-
-        let arcs = item_arcs
-            .iter()
-            .map(|arc| {
-                *finished
-                    .get(&arc.identity())
-                    .expect("nested arc should have been serialized first")
-            })
-            .collect();
-
-        self.store_data(PagableData {
-            data: item_data,
-            arcs,
-        })
-    }
-
     pub fn write_bytes<T: bytemuck::Pod>(&self, key: &str, data: T) {
         self.db
             .insert(key.as_bytes(), sled::IVec::from(bytemuck::bytes_of(&data)))
@@ -352,4 +289,31 @@ impl PagableStorage for SledBackedPagableStorage {
     fn session_context(&self) -> &Mutex<SessionContext> {
         &self.session_context
     }
+
+    /// Serialize `PagableData` into the on-disk byte format and insert into sled.
+    /// Returns the content-addressable `DataKey`.
+    fn store_data(&self, data: PagableData) -> DataKey {
+        let key = data.compute_key();
+        let db_key = bytemuck::bytes_of(&key);
+        if self
+            .db
+            .contains_key(db_key)
+            .expect("sled contains_key failed")
+        {
+            return key;
+        }
+
+        let bytes_size = 8 + 8 + data.data.len() + data.arcs.len() * 16;
+        let mut bytes = Vec::with_capacity(bytes_size);
+        bytes.extend_from_slice(&(data.data.len() as u64).to_le_bytes());
+        bytes.extend_from_slice(&(data.arcs.len() as u64).to_le_bytes());
+        bytes.extend_from_slice(&data.data);
+        bytes.extend_from_slice(bytemuck::cast_slice(&data.arcs));
+        assert_eq!(bytes.len(), bytes_size);
+
+        self.db
+            .insert(db_key, sled::IVec::from(bytes))
+            .expect("sled insert failed");
+        key
+    }
 }

pagable/src/storage/support.rs

@@ -27,7 +27,7 @@ pub struct SerializerForPaging<'a> {
 }
 
 impl<'a> SerializerForPaging<'a> {
-    pub(crate) fn new(session_context: &'a mut SessionContext) -> Self {
+    pub fn new(session_context: &'a mut SessionContext) -> Self {
         Self {
             serde: postcard::Serializer {
                 output: crate::flavors::PagableVecFlavor::new(),

pagable/src/storage/traits.rs

@@ -9,13 +9,15 @@
  */
 
 use std::any::TypeId;
+use std::collections::HashMap;
 use std::sync::Mutex;
 
 use either::Either;
 
 use crate::arc_erase::ArcEraseDyn;
 use crate::storage::data::DataKey;
 use crate::storage::data::PagableData;
+use crate::storage::support::SerializerForPaging;
 use crate::traits::SessionContext;
 
 /// Trait for storage backends that can persist and retrieve paged-out data.
@@ -68,6 +70,93 @@ pub trait PagableStorage: Send + Sync + 'static {
     /// Access the session context for storing/retrieving layer-specific state
     /// during serialization and deserialization.
     fn session_context(&self) -> &Mutex<SessionContext>;
+
+    /// Stores a single content-addressable [`PagableData`] blob and returns its
+    /// [`DataKey`]. The key is derived from the data via
+    /// [`PagableData::compute_key`]; if the same data is stored twice the second
+    /// write is expected to be idempotent (or skipped).
+    fn store_data(&self, data: PagableData) -> DataKey;
+
+    /// Stores a previously-serialized item (and its transitively reachable arcs)
+    /// to storage and returns its content-addressable [`DataKey`].
+    ///
+    /// The caller is responsible for the initial serialization: lock
+    /// `session_context()`, build a [`SerializerForPaging`], serialize the value,
+    /// `.finish()` to obtain `(item_data, item_arcs)`, then pass them in here
+    /// along with the still-locked `&mut SessionContext` (this method uses it to
+    /// recursively serialize nested arcs).
+    ///
+    /// `finished` is a cache of arc identity → `DataKey` that the caller may share
+    /// across multiple `page_out_item` invocations to avoid re-serializing arcs
+    /// that were already paged out earlier in the same batch.
+    fn page_out_item(
+        &self,
+        item_data: Vec<u8>,
+        item_arcs: Vec<Box<dyn ArcEraseDyn>>,
+        finished: &mut HashMap<usize, DataKey>,
+        session_context: &mut SessionContext,
+    ) -> DataKey {
+        enum Task {
+            Start(Box<dyn ArcEraseDyn>),
+            Finish((Box<dyn ArcEraseDyn>, Vec<u8>, Vec<Box<dyn ArcEraseDyn>>)),
+        }
+
+        let mut tasks: Vec<Task> = item_arcs
+            .iter()
+            .map(|arc| Task::Start(arc.clone_dyn()))
+            .collect();
+
+        while let Some(task) = tasks.pop() {
+            match task {
+                Task::Start(v) => {
+                    if finished.contains_key(&v.identity()) {
+                        continue;
+                    }
+
+                    let mut serializer = SerializerForPaging::new(session_context);
+                    v.serialize(&mut serializer).unwrap();
+                    let (data, arcs) = serializer.finish();
+
+                    let subtasks: Vec<_> = arcs
+                        .iter()
+                        .filter(|arc| !finished.contains_key(&arc.identity()))
+                        .map(|arc| Task::Start(arc.clone_dyn()))
+                        .collect();
+
+                    tasks.push(Task::Finish((v, data, arcs)));
+                    tasks.extend(subtasks);
+                }
+                Task::Finish((arc, data, serialized_arcs)) => {
+                    let arcs: Vec<DataKey> = serialized_arcs
+                        .iter()
+                        .map(|arc| {
+                            *finished
+                                .get(&arc.identity())
+                                .expect("nested arc should have been serialized first")
+                        })
+                        .collect();
+
+                    let key = self.store_data(PagableData { data, arcs });
+                    finished.insert(arc.identity(), key);
+                    arc.set_data_key(key);
+                }
+            }
+        }
+
+        let arcs: Vec<DataKey> = item_arcs
+            .iter()
+            .map(|arc| {
+                *finished
+                    .get(&arc.identity())
+                    .expect("nested arc should have been serialized first")
+            })
+            .collect();
+
+        self.store_data(PagableData {
+            data: item_data,
+            arcs,
+        })
+    }
 }
 
 static_assertions::assert_obj_safe!(PagableStorage);

pagable/src/testing.rs

@@ -257,4 +257,9 @@ impl PagableStorage for EmptyPagableStorage {
     fn session_context(&self) -> &Mutex<SessionContext> {
         &self.session_context
     }
+
+    fn store_data(&self, data: PagableData) -> DataKey {
+        // no-op storage; just compute the key
+        data.compute_key()
+    }
 }