shim: add wait flag

Author: gabriele-0201

sugondat/shim/src/cli.rs

@@ -171,6 +171,25 @@ pub mod query {
         Hash([u8; 32]),
     }
 
+    #[derive(Debug, Args)]
+    pub struct BlockParams {
+        /// The block containing the blob to query.
+        ///
+        /// Possible values: ["best", number, hash]
+        ///
+        /// "best" is the highest finalized block.
+        ///
+        /// Hashes must be 32 bytes, hex-encoded, and prefixed with "0x".
+        #[arg(value_name = "BLOCK_REF")]
+        pub block_ref: BlockRef,
+
+        /// By default, if the block is not available (e.g. not yet produced), the shim will return immediately.
+        ///
+        /// By specifying this flag, the shim will wait until the block is available.
+        #[clap(long)]
+        pub wait: bool,
+    }
+
     impl std::fmt::Display for BlockRef {
         fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
             match *self {
@@ -221,22 +240,15 @@ pub mod query {
     pub mod blob {
         use clap::Args;
 
-        use super::{BlockRef, SugondatRpcParams};
+        use super::{BlockParams, SugondatRpcParams};
 
         #[derive(Debug, Args)]
         pub struct Params {
             #[clap(flatten)]
             pub rpc: SugondatRpcParams,
 
-            /// The block containing the blob to query.
-            ///
-            /// Possible values: ["best", number, hash]
-            ///
-            /// "best" is the highest finalized block.
-            ///
-            /// Hashes must be 32 bytes, hex-encoded, and prefixed with "0x".
-            #[arg(value_name = "BLOCK_REF")]
-            pub block: BlockRef,
+            #[clap(flatten)]
+            pub block: BlockParams,
 
             /// The index of the extrinsic (transaction) containing the blob.
             #[arg(value_name = "INDEX")]
@@ -254,22 +266,15 @@ pub mod query {
 
         use clap::Args;
 
-        use super::{BlockRef, SugondatRpcParams};
+        use super::{BlockParams, SugondatRpcParams};
 
         #[derive(Debug, Args)]
         pub struct Params {
             #[clap(flatten)]
             pub rpc: SugondatRpcParams,
 
-            /// The block to query information about.
-            ///
-            /// Possible values: ["best", number, hash]
-            ///
-            /// "best" is the highest finalized block.
-            ///
-            /// Hashes must be 32 bytes, hex-encoded, and prefixed with "0x".
-            #[arg(default_value_t = BlockRef::Best, value_name = "BLOCK_REF")]
-            pub block: BlockRef,
+            #[clap(flatten)]
+            pub block: BlockParams,
         }
     }
 

sugondat/shim/src/cmd/query/blob.rs

@@ -1,5 +1,5 @@
-use super::connect_rpc;
-use crate::cli::query::{blob::Params, BlockRef};
+use super::{connect_rpc, get_block_at};
+use crate::cli::query::blob::Params;
 
 use std::io::Write;
 
@@ -12,20 +12,7 @@ pub async fn run(params: Params) -> anyhow::Result<()> {
     } = params;
 
     let client = connect_rpc(rpc).await?;
-
-    let maybe_hash = match block {
-        BlockRef::Best => None,
-        BlockRef::Hash(h) => Some(h),
-        BlockRef::Number(n) => Some(
-            client
-                .block_hash(n)
-                .await?
-                .ok_or_else(|| anyhow::anyhow!("No block with number {}", n))?
-                .0,
-        ),
-    };
-
-    let block = client.get_block_at(maybe_hash).await?;
+    let block = get_block_at(&client, block).await?;
 
     let i = block
         .blobs

sugondat/shim/src/cmd/query/block.rs

@@ -1,24 +1,11 @@
-use super::connect_rpc;
-use crate::cli::query::{block::Params, BlockRef};
+use super::{connect_rpc, get_block_at};
+use crate::cli::query::block::Params;
 
 pub async fn run(params: Params) -> anyhow::Result<()> {
     let Params { rpc, block } = params;
 
     let client = connect_rpc(rpc).await?;
-
-    let maybe_hash = match block {
-        BlockRef::Best => None,
-        BlockRef::Hash(h) => Some(h),
-        BlockRef::Number(n) => Some(
-            client
-                .block_hash(n)
-                .await?
-                .ok_or_else(|| anyhow::anyhow!("No block with number {}", n))?
-                .0,
-        ),
-    };
-
-    let block = client.get_block_at(maybe_hash).await?;
+    let block = get_block_at(&client, block).await?;
 
     println!("Block: #{}", block.number);
     println!("  Hash: 0x{}", hex::encode(&block.hash[..]));

sugondat/shim/src/cmd/query/mod.rs

@@ -1,5 +1,5 @@
 use crate::{
-    cli::query::{Commands, Params},
+    cli::query::{BlockParams, BlockRef, Commands, Params},
     sugondat_rpc,
 };
 
@@ -16,6 +16,30 @@ pub async fn run(params: Params) -> anyhow::Result<()> {
     Ok(())
 }
 
+/// Given the BlockParams and the client to be used, try to fetch
+/// the corresponding block. It will wait until the block is avaiable if specified.
+pub async fn get_block_at(
+    client: &sugondat_rpc::Client,
+    block: BlockParams,
+) -> anyhow::Result<sugondat_rpc::Block> {
+    let maybe_hash = match block.block_ref {
+        BlockRef::Best => None,
+        BlockRef::Hash(h) => Some(h),
+        BlockRef::Number(n) => Some(match block.wait {
+            true => client.wait_block_hash(n).await,
+            false => client
+                .block_hash(n)
+                .await?
+                .ok_or_else(|| anyhow::anyhow!("No block with number {}", n))?,
+        }),
+    };
+
+    match block.wait {
+        true => client.wait_block_at(maybe_hash).await,
+        false => client.get_block_at(maybe_hash).await,
+    }
+}
+
 async fn connect_rpc(
     conn_params: crate::cli::SugondatRpcParams,
 ) -> anyhow::Result<sugondat_rpc::Client> {

sugondat/shim/src/dock/rollkit.rs

@@ -34,7 +34,7 @@ impl RollkitRPCServer for RollkitDock {
         );
         let namespace = parse_namespace(&namespace).map_err(|_| err::bad_namespace())?;
         let block_hash = self.client.wait_finalized_height(height).await;
-        let block = self.client.get_block_at(Some(block_hash)).await.unwrap();
+        let block = self.client.wait_block_at(Some(block_hash)).await.unwrap();
         let mut blobs = vec![];
         for blob in block.blobs {
             if blob.namespace == namespace {

sugondat/shim/src/dock/sovereign.rs

@@ -34,7 +34,7 @@ impl SovereignRPCServer for SovereignDock {
     ) -> Result<Block, ErrorObjectOwned> {
         info!("get_block({})", height);
         let block_hash = self.client.wait_finalized_height(height).await;
-        let block = self.client.get_block_at(Some(block_hash)).await.unwrap();
+        let block = self.client.wait_block_at(Some(block_hash)).await.unwrap();
         let proof = make_namespace_proof(&block, namespace);
         let blobs = block
             .blobs

sugondat/shim/src/sugondat_rpc/mod.rs

@@ -2,7 +2,7 @@ use std::{fmt, sync::Arc};
 
 use crate::key::Keypair;
 use anyhow::Context;
-use subxt::{config::Header as _, rpc_params, utils::H256};
+use subxt::{config::Header as _, error::BlockError, rpc_params, utils::H256};
 use sugondat_nmt::Namespace;
 use sugondat_subxt::{
     sugondat::runtime_types::pallet_sugondat_blobs::namespace_param::UnvalidatedNamespace, Header,
@@ -63,6 +63,7 @@ impl Client {
     /// the block hash of the block at the given height.
     #[tracing::instrument(level = Level::DEBUG, skip(self))]
     pub async fn wait_finalized_height(&self, height: u64) -> [u8; 32] {
+        tracing::info!("Waiting for block at height: {}", height);
         loop {
             let conn = self.connector.ensure_connected().await;
             match conn.finalized.wait_until_finalized(self, height).await {
@@ -79,7 +80,7 @@ impl Client {
     ///
     /// If there is no block at the given height, returns `None`.
     #[tracing::instrument(level = Level::DEBUG, skip(self))]
-    pub async fn block_hash(&self, height: u64) -> anyhow::Result<Option<H256>> {
+    pub async fn block_hash(&self, height: u64) -> anyhow::Result<Option<[u8; 32]>> {
         loop {
             let conn = self.connector.ensure_connected().await;
             let block_hash: Option<H256> = match conn
@@ -102,43 +103,75 @@ impl Client {
                     // at the given height.
                     Ok(None)
                 }
-                Some(block_hash) => Ok(Some(block_hash)),
+                Some(block_hash) => Ok(Some(block_hash.into())),
             };
         }
     }
 
-    /// Returns the header and the body of the block with the given hash, automatically retrying
-    /// until it succeeds. `None` indicates the best block.
+    /// Returns the block hash of the block at the given height,
+    /// automatically retrying until it succeeds.
+    /// `None` indicates the best block.
+    pub async fn wait_block_hash(&self, height: u64) -> [u8; 32] {
+        loop {
+            match self.block_hash(height).await {
+                Ok(Some(res)) => break res,
+                Ok(None) => {
+                    tracing::info!("Block hash not available yet");
+                    tokio::time::sleep(std::time::Duration::from_secs(1)).await;
+                }
+                Err(e) => {
+                    // any other error is treated as a failure, resulting in the connection being reset
+                    tracing::error!(?e, "failed to query block hash");
+                    self.connector.reset().await;
+                }
+            }
+        }
+    }
+
+    /// Returns the header and body of the block with the given hash.
+    /// If it is not available, it returns the reason of the failure.
+    /// `None` indicates the best block.
     async fn get_header_and_extrinsics(
         &self,
         block_hash: Option<[u8; 32]>,
-    ) -> anyhow::Result<(Header, Vec<sugondat_subxt::ExtrinsicDetails>)> {
+    ) -> Result<(Header, Vec<sugondat_subxt::ExtrinsicDetails>), subxt::error::Error> {
         let block_hash = block_hash.map(H256::from);
+
+        let conn = self.connector.ensure_connected().await;
+        let res = match block_hash {
+            Some(h) => conn.subxt.blocks().at(h).await,
+            None => conn.subxt.blocks().at_latest().await,
+        }?;
+
+        let header = res.header();
+        let body = res
+            .extrinsics()
+            .await?
+            .iter()
+            .collect::<Result<Vec<_>, _>>()?;
+        Ok((header.clone(), body))
+    }
+
+    /// Returns the header and the body of the block with the given hash,
+    /// automatically retrying until it succeeds.
+    /// `None` indicates the best block.
+    async fn wait_header_and_extrinsics(
+        &self,
+        block_hash: Option<[u8; 32]>,
+    ) -> (Header, Vec<sugondat_subxt::ExtrinsicDetails>) {
         loop {
-            let conn = self.connector.ensure_connected().await;
-            let res = match block_hash {
-                Some(h) => conn.subxt.blocks().at(h).await,
-                None => conn.subxt.blocks().at_latest().await,
-            };
-            let err = match res {
-                Ok(it) => {
-                    let header = it.header();
-                    let body = match it.extrinsics().await {
-                        Ok(it) => it,
-                        Err(err) => {
-                            tracing::error!(?err, "failed to query block");
-                            self.connector.reset().await;
-                            continue;
-                        }
-                    }
-                    .iter()
-                    .collect::<Result<Vec<_>, _>>()?;
-                    return Ok((header.clone(), body));
+            match self.get_header_and_extrinsics(block_hash).await {
+                Ok(res) => break res,
+                Err(subxt::Error::Block(BlockError::NotFound(_hash))) => {
+                    tracing::info!("Block not available yet");
+                    tokio::time::sleep(std::time::Duration::from_secs(1)).await;
                 }
-                Err(err) => err,
-            };
-            tracing::error!(?err, "failed to query block");
-            self.connector.reset().await;
+                Err(e) => {
+                    // any other error is treated as a failure, resulting in the connection being reset
+                    tracing::error!(?e, "failed to query block");
+                    self.connector.reset().await;
+                }
+            }
         }
     }
 
@@ -148,19 +181,16 @@ impl Client {
     /// `None` indicates that the best block should be used.
     #[tracing::instrument(level = Level::DEBUG, skip(self))]
     pub async fn get_block_at(&self, block_hash: Option<[u8; 32]>) -> anyhow::Result<Block> {
-        let (header, extrinsics) = self.get_header_and_extrinsics(block_hash).await?;
-        let tree_root = tree_root(&header).ok_or_else(err::no_tree_root)?;
-        let timestamp = extract_timestamp(&extrinsics)?;
-        let blobs = extract_blobs(extrinsics);
-        tracing::debug!(?blobs, "found {} blobs in block", blobs.len());
-        Ok(Block {
-            number: header.number as u64,
-            hash: header.hash().0,
-            parent_hash: header.parent_hash.0,
-            tree_root,
-            timestamp,
-            blobs,
-        })
+        Block::from_header_and_extrinsics(self.get_header_and_extrinsics(block_hash).await?)
+    }
+
+    /// Returns the data of the block identified by the given block hash.
+    /// It retries until the block is present
+    ///
+    /// `None` indicates that the best block should be used.
+    #[tracing::instrument(level = Level::DEBUG, skip(self))]
+    pub async fn wait_block_at(&self, block_hash: Option<[u8; 32]>) -> anyhow::Result<Block> {
+        Block::from_header_and_extrinsics(self.wait_header_and_extrinsics(block_hash).await)
     }
 
     /// Submit a blob with the given namespace and signed with the given key. The block is submitted
@@ -319,7 +349,7 @@ impl FinalizedHeadWatcher {
                     // TODO: throttle the retries. // TODO: return an error
                     continue;
                 };
-                break Some(block_hash.0);
+                break Some(block_hash);
             }
         }
     }
@@ -347,6 +377,26 @@ pub struct Block {
     pub blobs: Vec<Blob>,
 }
 
+impl Block {
+    fn from_header_and_extrinsics(
+        value: (Header, Vec<sugondat_subxt::ExtrinsicDetails>),
+    ) -> anyhow::Result<Self> {
+        let (header, extrinsics) = value;
+        let tree_root = tree_root(&header).ok_or_else(err::no_tree_root)?;
+        let timestamp = extract_timestamp(&extrinsics)?;
+        let blobs = extract_blobs(extrinsics);
+        tracing::debug!(?blobs, "found {} blobs in block", blobs.len());
+        Ok(Block {
+            number: header.number as u64,
+            hash: header.hash().0,
+            parent_hash: header.parent_hash.0,
+            tree_root,
+            timestamp,
+            blobs,
+        })
+    }
+}
+
 /// Represents a blob in a sugondat block.
 pub struct Blob {
     pub extrinsic_index: u32,