Make blocks queryable by height (#422)

Cargo.lock

@@ -829,6 +829,12 @@ dependencies = [
  "regex-automata",
 ]
 
+[[package]]
+name = "matches"
+version = "0.1.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7ffc5c5338469d4d3ea17d269fa8ea3512ad247247c30bd2df69e68309ed0a08"
+
 [[package]]
 name = "maybe-uninit"
 version = "2.0.0"
@@ -1106,7 +1112,7 @@ dependencies = [
  "libc",
  "redox_syscall",
  "rustc_version",
- "smallvec",
+ "smallvec 0.6.13",
  "winapi 0.3.8",
 ]
 
@@ -1552,6 +1558,12 @@ dependencies = [
  "maybe-uninit",
 ]
 
+[[package]]
+name = "smallvec"
+version = "1.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c7cb5678e1615754284ec264d9bb5b4c27d2018577fd90ac0ceb578591ed5ee4"
+
 [[package]]
 name = "socket2"
 version = "0.3.12"
@@ -1564,6 +1576,19 @@ dependencies = [
  "winapi 0.3.8",
 ]
 
+[[package]]
+name = "spandoc"
+version = "0.1.3"
+source = "git+https://github.com/yaahc/spandoc.git#554358be632b156a6f0af963b0b244e2665b4767"
+dependencies = [
+ "matches",
+ "proc-macro2 1.0.9",
+ "quote 1.0.3",
+ "syn 1.0.17",
+ "tracing",
+ "tracing-futures",
+]
+
 [[package]]
 name = "stable_deref_trait"
 version = "1.1.1"
@@ -1931,6 +1956,16 @@ dependencies = [
  "tracing-core",
 ]
 
+[[package]]
+name = "tracing-serde"
+version = "0.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b6ccba2f8f16e0ed268fc765d9b7ff22e965e7185d32f8f1ec8294fe17d86e79"
+dependencies = [
+ "serde",
+ "tracing-core",
+]
+
 [[package]]
 name = "tracing-subscriber"
 version = "0.1.6"
@@ -1943,7 +1978,7 @@ dependencies = [
  "matchers",
  "owning_ref",
  "regex",
- "smallvec",
+ "smallvec 0.6.13",
  "tracing-core",
  "tracing-log",
 ]
@@ -1954,8 +1989,18 @@ version = "0.2.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "1d53c40489aa69c9aed21ff483f26886ca8403df33bdc2d2f87c60c1617826d2"
 dependencies = [
+ "ansi_term",
+ "chrono",
+ "lazy_static",
+ "matchers",
+ "regex",
+ "serde",
+ "serde_json",
  "sharded-slab",
+ "smallvec 1.4.0",
  "tracing-core",
+ "tracing-log",
+ "tracing-serde",
 ]
 
 [[package]]
@@ -2160,8 +2205,13 @@ dependencies = [
  "futures",
  "hex",
  "lazy_static",
+ "spandoc",
  "tokio",
  "tower",
+ "tracing",
+ "tracing-error",
+ "tracing-futures",
+ "tracing-subscriber 0.2.5",
  "zebra-chain",
  "zebra-test-vectors",
 ]

zebra-state/Cargo.toml

@@ -20,3 +20,8 @@ color-eyre = "0.3.4"
 eyre = "0.4.2"
 tokio = { version = "0.2.21", features = ["full"] }
 zebra-test-vectors = { path = "../zebra-test-vectors/" }
+spandoc = { git = "https://github.com/yaahc/spandoc.git" }
+tracing = "0.1.15"
+tracing-futures = "0.2.4"
+tracing-error = "0.1.2"
+tracing-subscriber = "0.2.5"

zebra-state/src/in_memory.rs

@@ -0,0 +1,67 @@
+use super::{Request, Response};
+use futures::prelude::*;
+use std::{
+    error::Error,
+    future::Future,
+    pin::Pin,
+    task::{Context, Poll},
+};
+use tower::{buffer::Buffer, Service};
+
+mod block_index;
+
+#[derive(Default)]
+struct ZebraState {
+    index: block_index::BlockIndex,
+}
+
+impl Service<Request> for ZebraState {
+    type Response = Response;
+    type Error = Box<dyn Error + Send + Sync + 'static>;
+    type Future =
+        Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send + 'static>>;
+
+    fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        Poll::Ready(Ok(()))
+    }
+
+    fn call(&mut self, req: Request) -> Self::Future {
+        match req {
+            Request::AddBlock { block } => {
+                let result = self.index.insert(block).map(|_| Response::Added);
+
+                async { result }.boxed()
+            }
+            Request::GetBlock { hash } => {
+                let result = self
+                    .index
+                    .get(hash)
+                    .map(|block| Response::Block { block })
+                    .ok_or_else(|| "block could not be found".into());
+
+                async move { result }.boxed()
+            }
+            Request::GetTip => {
+                let result = self
+                    .index
+                    .get_tip()
+                    .map(|block| block.as_ref().into())
+                    .map(|hash| Response::Tip { hash })
+                    .ok_or_else(|| "zebra-state contains no blocks".into());
+
+                async move { result }.boxed()
+            }
+        }
+    }
+}
+
+pub fn init() -> impl Service<
+    Request,
+    Response = Response,
+    Error = Box<dyn Error + Send + Sync + 'static>,
+    Future = impl Future<Output = Result<Response, Box<dyn Error + Send + Sync + 'static>>>,
+> + Send
+       + Clone
+       + 'static {
+    Buffer::new(ZebraState::default(), 1)
+}

zebra-state/src/in_memory/block_index.rs

@@ -0,0 +1,67 @@
+use std::{
+    collections::{btree_map::Entry, BTreeMap, HashMap},
+    error::Error,
+    sync::Arc,
+};
+use zebra_chain::{
+    block::{Block, BlockHeaderHash},
+    types::BlockHeight,
+};
+#[derive(Default)]
+pub(super) struct BlockIndex {
+    by_hash: HashMap<BlockHeaderHash, Arc<Block>>,
+    by_height: BTreeMap<BlockHeight, Arc<Block>>,
+}
+
+impl BlockIndex {
+    pub(super) fn insert(
+        &mut self,
+        block: impl Into<Arc<Block>>,
+    ) -> Result<(), Box<dyn Error + Send + Sync + 'static>> {
+        let block = block.into();
+        let hash = block.as_ref().into();
+        let height = block.coinbase_height().unwrap();
+
+        match self.by_height.entry(height) {
+            Entry::Vacant(entry) => {
+                let _ = entry.insert(block.clone());
+                let _ = self.by_hash.insert(hash, block);
+                Ok(())
+            }
+            Entry::Occupied(_) => Err("forks in the chain aren't supported yet")?,
+        }
+    }
+
+    pub(super) fn get(&mut self, query: impl Into<BlockQuery>) -> Option<Arc<Block>> {
+        match query.into() {
+            BlockQuery::ByHash(hash) => self.by_hash.get(&hash),
+            BlockQuery::ByHeight(height) => self.by_height.get(&height),
+        }
+        .cloned()
+    }
+
+    pub(super) fn get_tip(&self) -> Option<Arc<Block>> {
+        self.by_height
+            .iter()
+            .next_back()
+            .map(|(_key, value)| value)
+            .cloned()
+    }
+}
+
+pub(super) enum BlockQuery {
+    ByHash(BlockHeaderHash),
+    ByHeight(BlockHeight),
+}
+
+impl From<BlockHeaderHash> for BlockQuery {
+    fn from(hash: BlockHeaderHash) -> Self {
+        Self::ByHash(hash)
+    }
+}
+
+impl From<BlockHeight> for BlockQuery {
+    fn from(height: BlockHeight) -> Self {
+        Self::ByHeight(height)
+    }
+}

zebra-state/src/lib.rs

@@ -1,81 +1,24 @@
 #![doc(html_logo_url = "https://www.zfnd.org/images/zebra-icon.png")]
 #![doc(html_root_url = "https://doc.zebra.zfnd.org/zebra_state")]
-use futures::prelude::*;
+#![allow(clippy::try_err)]
-use std::{
+use std::sync::Arc;
-    collections::HashMap,
-    error::Error,
-    future::Future,
-    pin::Pin,
-    sync::Arc,
-    task::{Context, Poll},
-};
-use tower::{buffer::Buffer, Service};
 use zebra_chain::block::{Block, BlockHeaderHash};
 
+pub mod in_memory;
+
 #[derive(Debug)]
 pub enum Request {
     // TODO(jlusby): deprecate in the future based on our validation story
     AddBlock { block: Arc<Block> },
     GetBlock { hash: BlockHeaderHash },
+    GetTip,
 }
 
 #[derive(Debug)]
 pub enum Response {
     Added,
     Block { block: Arc<Block> },
-}
+    Tip { hash: BlockHeaderHash },
-
-pub mod in_memory {
-    use super::*;
-    use std::error::Error;
-
-    pub fn init() -> impl Service<
-        Request,
-        Response = Response,
-        Error = Box<dyn Error + Send + Sync + 'static>,
-        Future = impl Future<Output = Result<Response, Box<dyn Error + Send + Sync + 'static>>>,
-    > + Send
-           + Clone
-           + 'static {
-        Buffer::new(ZebraState::default(), 1)
-    }
-}
-
-#[derive(Default)]
-struct ZebraState {
-    blocks: HashMap<BlockHeaderHash, Arc<Block>>,
-}
-
-impl Service<Request> for ZebraState {
-    type Response = Response;
-    type Error = Box<dyn Error + Send + Sync + 'static>;
-    type Future =
-        Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send + 'static>>;
-
-    fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
-        Poll::Ready(Ok(()))
-    }
-
-    fn call(&mut self, req: Request) -> Self::Future {
-        match req {
-            Request::AddBlock { block } => {
-                let hash = block.as_ref().into();
-                self.blocks.insert(hash, block);
-
-                async { Ok(Response::Added) }.boxed()
-            }
-            Request::GetBlock { hash } => {
-                let result = self
-                    .blocks
-                    .get(&hash)
-                    .cloned()
-                    .map(|block| Response::Block { block })
-                    .ok_or_else(|| "block could not be found".into());
-
-                async move { result }.boxed()
-            }
-        }
-    }
 }
 
 #[cfg(test)]
@@ -83,8 +26,26 @@ mod tests {
     use super::*;
     use color_eyre::Report;
     use eyre::{bail, ensure, eyre};
+    use tower::Service;
     use zebra_chain::serialization::ZcashDeserialize;
 
+    fn install_tracing() {
+        use tracing_error::ErrorLayer;
+        use tracing_subscriber::prelude::*;
+        use tracing_subscriber::{fmt, EnvFilter};
+
+        let fmt_layer = fmt::layer().with_target(false);
+        let filter_layer = EnvFilter::try_from_default_env()
+            .or_else(|_| EnvFilter::try_new("info"))
+            .unwrap();
+
+        tracing_subscriber::registry()
+            .with(filter_layer)
+            .with(fmt_layer)
+            .with(ErrorLayer::default())
+            .init();
+    }
+
     #[tokio::test]
     async fn round_trip() -> Result<(), Report> {
         let block: Arc<_> =
@@ -120,4 +81,55 @@ mod tests {
 
         Ok(())
     }
+
+    #[tokio::test]
+    #[spandoc::spandoc]
+    async fn get_tip() -> Result<(), Report> {
+        install_tracing();
+
+        let block0: Arc<_> =
+            Block::zcash_deserialize(&zebra_test_vectors::BLOCK_MAINNET_GENESIS_BYTES[..])?.into();
+        let block1: Arc<_> =
+            Block::zcash_deserialize(&zebra_test_vectors::BLOCK_MAINNET_1_BYTES[..])?.into();
+
+        let expected_hash: BlockHeaderHash = block1.as_ref().into();
+
+        let mut service = in_memory::init();
+
+        /// insert the higher block first
+        let response = service
+            .call(Request::AddBlock { block: block1 })
+            .await
+            .map_err(|e| eyre!(e))?;
+
+        ensure!(
+            matches!(response, Response::Added),
+            "unexpected response kind: {:?}",
+            response
+        );
+
+        /// genesis block second
+        let response = service
+            .call(Request::AddBlock {
+                block: block0.clone(),
+            })
+            .await
+            .map_err(|e| eyre!(e))?;
+
+        ensure!(
+            matches!(response, Response::Added),
+            "unexpected response kind: {:?}",
+            response
+        );
+
+        let block_response = service.call(Request::GetTip).await.map_err(|e| eyre!(e))?;
+
+        /// assert that the higher block is returned as the tip even tho it was least recently inserted
+        match block_response {
+            Response::Tip { hash } => assert_eq!(expected_hash, hash),
+            _ => bail!("unexpected response kind: {:?}", block_response),
+        }
+
+        Ok(())
+    }
 }