Merge pull request #2979 from autonomys/obj-map-rpc

Add RPC subscription for piece object mappings

Author: teor2345

crates/sc-consensus-subspace-rpc/src/lib.rs

@@ -20,7 +20,7 @@
 #![feature(try_blocks)]
 
 use futures::channel::mpsc;
-use futures::{future, FutureExt, StreamExt};
+use futures::{future, stream, FutureExt, StreamExt};
 use jsonrpsee::core::async_trait;
 use jsonrpsee::proc_macros::rpc;
 use jsonrpsee::types::{ErrorObject, ErrorObjectOwned};
@@ -58,6 +58,7 @@ use std::sync::{Arc, Weak};
 use std::time::Duration;
 use subspace_archiving::archiver::NewArchivedSegment;
 use subspace_core_primitives::crypto::kzg::Kzg;
+use subspace_core_primitives::objects::GlobalObject;
 use subspace_core_primitives::{
     BlockHash, HistorySize, Piece, PieceIndex, PublicKey, SegmentHeader, SegmentIndex, SlotNumber,
     Solution,
@@ -155,6 +156,16 @@ pub trait SubspaceRpcApi {
 
     #[method(name = "subspace_lastSegmentHeaders")]
     async fn last_segment_headers(&self, limit: u64) -> Result<Vec<Option<SegmentHeader>>, Error>;
+
+    /// Block/transaction archived object mappings subscription
+    #[subscription(
+        name = "subspace_subscribeArchivedObjectMappings" => "subspace_archived_object_mappings",
+        unsubscribe = "subspace_unsubscribeArchivedObjectMappings",
+        item = GlobalObject,
+    )]
+    fn subscribe_archived_object_mappings(&self);
+
+    // TODO: add a method for recent/any object mappings based on a list of IDs, piece indexes, or segment indexes
 }
 
 #[derive(Default)]
@@ -703,6 +714,7 @@ where
         Ok(())
     }
 
+    // Note: this RPC uses the cached archived segment, which is only updated by archived segments subscriptions
     fn piece(&self, requested_piece_index: PieceIndex) -> Result<Option<Piece>, Error> {
         self.deny_unsafe.check_if_safe()?;
 
@@ -807,4 +819,30 @@ where
 
         Ok(last_segment_headers)
     }
+
+    // TODO:
+    // - the number of object mappings in each segment can be very large (hundreds or thousands).
+    //   To avoid RPC connection failures, limit the number of mappings returned in each response,
+    //   or the number of in-flight responses.
+    fn subscribe_archived_object_mappings(&self, pending: PendingSubscriptionSink) {
+        // The genesis segment isn't included in this stream. In other methods we recreate is as the first segment,
+        // but there aren't any mappings in it, so we don't need to recreate it as part of this subscription.
+
+        let stream = self
+            .archived_segment_notification_stream
+            .subscribe()
+            .flat_map(|archived_segment_notification| {
+                let objects = archived_segment_notification
+                    .archived_segment
+                    .global_object_mappings();
+
+                stream::iter(objects)
+            });
+
+        self.subscription_executor.spawn(
+            "subspace-archived-object-mappings-subscription",
+            Some("rpc"),
+            pipe_from_stream(pending, stream).boxed(),
+        );
+    }
 }

crates/subspace-archiving/src/archiver.rs

@@ -34,7 +34,7 @@ use rayon::prelude::*;
 use subspace_core_primitives::crypto::kzg::{Commitment, Kzg, Witness};
 use subspace_core_primitives::crypto::{blake3_254_hash_to_scalar, Scalar};
 use subspace_core_primitives::objects::{
-    BlockObject, BlockObjectMapping, PieceObject, PieceObjectMapping,
+    BlockObject, BlockObjectMapping, GlobalObject, PieceObject, PieceObjectMapping,
 };
 use subspace_core_primitives::{
     ArchivedBlockProgress, ArchivedHistorySegment, Blake3Hash, BlockNumber, LastArchivedBlock,
@@ -184,6 +184,34 @@ pub struct NewArchivedSegment {
     pub object_mapping: Vec<PieceObjectMapping>,
 }
 
+impl NewArchivedSegment {
+    /// Returns all the object mappings in this archived segment as a lazy iterator.
+    pub fn global_object_mappings(&self) -> impl Iterator<Item = GlobalObject> + 'static {
+        // Save memory by only returning the necessary parts of NewArchivedSegment
+        let object_mapping = self.object_mapping.clone();
+        let piece_indexes = self
+            .segment_header
+            .segment_index()
+            .segment_piece_indexes_source_first();
+
+        // Iterate through the object mapping vector for each piece
+        object_mapping
+            .into_iter()
+            .zip(piece_indexes)
+            .flat_map(|(piece_mappings, piece_index)| {
+                // And then through each individual object mapping in the piece
+                piece_mappings
+                    .objects
+                    .into_iter()
+                    .map(move |piece_object| GlobalObject::V0 {
+                        piece_index,
+                        offset: piece_object.offset(),
+                        hash: piece_object.hash(),
+                    })
+            })
+    }
+}
+
 /// Archiver instantiation error
 #[derive(Debug, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
 #[cfg_attr(feature = "thiserror", derive(thiserror::Error))]

crates/subspace-core-primitives/src/objects.rs

@@ -83,7 +83,7 @@ pub struct BlockObjectMapping {
     pub objects: Vec<BlockObject>,
 }
 
-/// Object stored inside of the block
+/// Object stored inside of the piece
 #[derive(Debug, Copy, Clone, PartialEq, Eq, Ord, PartialOrd, Hash, Encode, Decode, TypeInfo)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "serde", serde(rename_all = "camelCase"))]
@@ -95,7 +95,7 @@ pub enum PieceObject {
         /// Object hash
         hash: Blake3Hash,
         // TODO: This is a raw record offset, not a regular one
-        /// Offset of the object
+        /// Offset of the object in that piece
         offset: u32,
     },
 }
@@ -121,11 +121,11 @@ impl PieceObject {
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "serde", serde(rename_all = "camelCase"))]
 pub struct PieceObjectMapping {
-    /// Objects stored inside of the block
+    /// Objects stored inside of the piece
     pub objects: Vec<PieceObject>,
 }
 
-/// Object stored inside in the history of the blockchain
+/// Object stored in the history of the blockchain
 #[derive(Debug, Copy, Clone, PartialEq, Eq, Ord, PartialOrd, Hash, Encode, Decode, TypeInfo)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "serde", serde(rename_all = "camelCase"))]
@@ -134,10 +134,13 @@ pub enum GlobalObject {
     #[codec(index = 0)]
     #[cfg_attr(feature = "serde", serde(rename_all = "camelCase"))]
     V0 {
-        /// Piece index where object is contained (at least its beginning, might not fit fully)
+        /// Piece index where object is contained (at least its beginning, might not fit fully).
+        /// The index and offset must be first, so that the Ord derives are valid.
         piece_index: PieceIndex,
-        /// Offset of the object
+        /// Offset of the object in that piece
         offset: u32,
+        /// Object hash
+        hash: Blake3Hash,
     },
 }
 
@@ -155,4 +158,21 @@ impl GlobalObject {
             Self::V0 { offset, .. } => *offset,
         }
     }
+
+    /// Object hash
+    pub fn hash(&self) -> Blake3Hash {
+        match self {
+            Self::V0 { hash, .. } => *hash,
+        }
+    }
+}
+
+/// Mapping of objects stored in the history of the blockchain
+#[derive(Debug, Default, Clone, PartialEq, Eq, Ord, PartialOrd, Hash, Encode, Decode, TypeInfo)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "serde", serde(rename_all = "camelCase"))]
+pub struct GlobalObjectMapping {
+    /// Objects stored in the history of the blockchain.
+    /// Mappings are ordered by the piece index and offset of the first GlobalObject in objects.
+    pub objects: Vec<GlobalObject>,
 }