Repository: dbpunk-labs/db3 (language: rust)

Issue:
LightStore: add block subscribe and block fetch api
## Motivations

<!--
If your feature request is related to a problem, please describe it.
-->
Add block subscribe for light storage node and support API to fetch block/mutation based on block height
- Would you like to implement this feature? [y/n]

## Solution

<!-- Describe the solution you'd like. -->

## Alternatives

<!-- Describe any alternative solutions or features you've considered. -->

## Additional context

<!-- Add any other context or screenshots about the feature request here. -->



Discussion / hints:

Fix patch:
diff --git a/src/node/src/command.rs b/src/node/src/command.rs
index 310ef055..1ed9f429 100644
--- a/src/node/src/command.rs
+++ b/src/node/src/command.rs
@@ -32,6 +32,9 @@ use db3_proto::db3_indexer_proto::indexer_node_server::IndexerNodeServer;
 use db3_proto::db3_node_proto::storage_node_client::StorageNodeClient;
 use db3_proto::db3_node_proto::storage_node_server::StorageNodeServer;
 use db3_proto::db3_storage_proto::storage_node_server::StorageNodeServer as StorageNodeV2Server;
+use db3_proto::db3_storage_proto::{
+    EventMessage as EventMessageV2, Subscription as SubscriptionV2,
+};
 use db3_sdk::mutation_sdk::MutationSDK;
 use db3_sdk::store_sdk::StoreSDK;
 use db3_storage::db_store_v2::DBStoreV2Config;
@@ -474,6 +477,11 @@ impl DB3Command {
             scan_max_limit: 1000,
         };
 
+        let (sender, receiver) = tokio::sync::mpsc::channel::<(
+            DB3Address,
+            SubscriptionV2,
+            Sender<std::result::Result<EventMessageV2, Status>>,
+        )>(1024);
         let config = StorageNodeV2Config {
             store_config,
             state_config,
@@ -482,13 +490,14 @@ impl DB3Command {
             network_id,
             block_interval,
         };
-        let storage_node = StorageNodeV2Impl::new(config).unwrap();
+        let storage_node = StorageNodeV2Impl::new(config, sender).unwrap();
         info!(
             "start db3 store node on public addr {} and network {}",
             addr, network_id
         );
         std::fs::create_dir_all(rollup_data_path).unwrap();
-        storage_node.start_to_produce_block();
+        storage_node.keep_subscription(receiver).await.unwrap();
+        storage_node.start_to_produce_block().await;
         storage_node.start_to_rollup().await;
         let cors_layer = CorsLayer::new()
             .allow_methods([Method::GET, Method::POST, Method::OPTIONS])
@@ -543,6 +552,7 @@ impl DB3Command {
                 .await
                 .unwrap();
         }
+        info!("db3 storage node exit");
     }
 
     ///
diff --git a/src/node/src/storage_node_light_impl.rs b/src/node/src/storage_node_light_impl.rs
index afb4df62..5ab6ad34 100644
--- a/src/node/src/storage_node_light_impl.rs
+++ b/src/node/src/storage_node_light_impl.rs
@@ -18,31 +18,41 @@
 use crate::mutation_utils::MutationUtil;
 use crate::rollup_executor::{RollupExecutor, RollupExecutorConfig};
 use db3_crypto::db3_address::DB3Address;
+use db3_crypto::db3_verifier::DB3Verifier;
 use db3_crypto::id::TxId;
 use db3_error::Result;
 use db3_proto::db3_mutation_v2_proto::{
     mutation::body_wrapper::Body, MutationAction, MutationRollupStatus,
 };
+use db3_proto::db3_storage_proto::event_message::Event as EventV2;
 use db3_proto::db3_storage_proto::{
     storage_node_server::StorageNode, ExtraItem, GetCollectionOfDatabaseRequest,
     GetCollectionOfDatabaseResponse, GetDatabaseOfOwnerRequest, GetDatabaseOfOwnerResponse,
     GetMutationBodyRequest, GetMutationBodyResponse, GetMutationHeaderRequest,
     GetMutationHeaderResponse, GetNonceRequest, GetNonceResponse, ScanMutationHeaderRequest,
     ScanMutationHeaderResponse, ScanRollupRecordRequest, ScanRollupRecordResponse,
-    SendMutationRequest, SendMutationResponse,
+    SendMutationRequest, SendMutationResponse, SubscribeRequest,
+};
+use db3_proto::db3_storage_proto::{
+    BlockEvent as BlockEventV2, EventMessage as EventMessageV2, EventType as EventTypeV2,
+    Subscription as SubscriptionV2,
 };
 use db3_storage::db_store_v2::{DBStoreV2, DBStoreV2Config};
 use db3_storage::mutation_store::{MutationStore, MutationStoreConfig};
 use db3_storage::state_store::{StateStore, StateStoreConfig};
+use prost::Message;
+use std::collections::{BTreeMap, HashSet};
 use std::sync::atomic::AtomicBool;
 use std::sync::atomic::Ordering;
 use std::sync::Arc;
-use std::time::Duration;
+use tokio::sync::broadcast;
+use tokio::sync::broadcast::Sender as BroadcastSender;
+use tokio::sync::mpsc::{Receiver, Sender};
 use tokio::task;
 use tokio::time::{sleep, Duration as TokioDuration};
+use tokio_stream::wrappers::ReceiverStream;
 use tonic::{Request, Response, Status};
-use tracing::{info, warn};
-
+use tracing::{debug, info, warn};
 pub struct StorageNodeV2Config {
     pub store_config: MutationStoreConfig,
     pub state_config: StateStoreConfig,
@@ -58,37 +68,77 @@ pub struct StorageNodeV2Impl {
     config: StorageNodeV2Config,
     running: Arc<AtomicBool>,
     db_store: DBStoreV2,
+    sender: Sender<(
+        DB3Address,
+        SubscriptionV2,
+        Sender<std::result::Result<EventMessageV2, Status>>,
+    )>,
+    broadcast_sender: BroadcastSender<EventMessageV2>,
 }
 
 impl StorageNodeV2Impl {
-    pub fn new(config: StorageNodeV2Config) -> Result<Self> {
+    pub fn new(
+        config: StorageNodeV2Config,
+        sender: Sender<(
+            DB3Address,
+            SubscriptionV2,
+            Sender<std::result::Result<EventMessageV2, Status>>,
+        )>,
+    ) -> Result<Self> {
         let storage = MutationStore::new(config.store_config.clone())?;
         let state_store = StateStore::new(config.state_config.clone())?;
         let db_store = DBStoreV2::new(config.db_store_config.clone())?;
+        let (broadcast_sender, _) = broadcast::channel(1024);
         Ok(Self {
             storage,
             state_store,
             config,
             running: Arc::new(AtomicBool::new(true)),
             db_store,
+            sender,
+            broadcast_sender,
         })
     }
 
-    pub fn start_to_produce_block(&self) {
+    pub async fn start_to_produce_block(&self) {
         let local_running = self.running.clone();
         let local_storage = self.storage.clone();
         let local_block_interval = self.config.block_interval;
-        task::spawn_blocking(move || {
+        let local_event_sender = self.broadcast_sender.clone();
+        task::spawn(async move {
             info!("start the block producer thread");
             while local_running.load(Ordering::Relaxed) {
-                std::thread::sleep(Duration::from_millis(local_block_interval));
+                sleep(TokioDuration::from_millis(local_block_interval)).await;
+                info!(
+                    "produce block {}",
+                    local_storage.get_current_block().unwrap_or(0)
+                );
                 match local_storage.increase_block() {
-                    Ok(()) => {}
+                    Ok((block_id, mutation_count)) => {
+                        // sender block event
+                        let e = BlockEventV2 {
+                            block_id,
+                            mutation_count,
+                        };
+                        let msg = EventMessageV2 {
+                            r#type: EventTypeV2::Block as i32,
+                            event: Some(EventV2::BlockEvent(e)),
+                        };
+                        match local_event_sender.send(msg) {
+                            Ok(_) => {
+                                info!("broadcast block event {}, {}", block_id, mutation_count);
+                            }
+                            Err(e) => {
+                                warn!("the broadcast channel error for {:?}", e);
+                            }
+                        }
+                    }
                     Err(e) => {
                         warn!("fail to produce block for error {e}");
                     }
                 }
             }
+            info!("exit the block producer thread");
         });
     }
 
@@ -110,12 +160,124 @@ impl StorageNodeV2Impl {
                     }
                 }
             }
+            info!("exit the rollup thread");
         });
     }
+    pub async fn keep_subscription(
+        &self,
+        mut receiver: Receiver<(
+            DB3Address,
+            SubscriptionV2,
+            Sender<std::result::Result<EventMessageV2, Status>>,
+        )>,
+    ) -> std::result::Result<(), Status> {
+        info!("start to keep subscription");
+        let local_running = self.running.clone();
+        let local_broadcast_sender = self.broadcast_sender.clone();
+
+        tokio::spawn(async move {
+            info!("listen to subscription update event and event message broadcaster");
+            while local_running.load(Ordering::Relaxed) {
+                info!("keep subscription loop");
+                let mut subscribers: BTreeMap<
+                    DB3Address,
+                    (
+                        Sender<std::result::Result<EventMessageV2, Status>>,
+                        SubscriptionV2,
+                    ),
+                > = BTreeMap::new();
+                let mut to_be_removed: HashSet<DB3Address> = HashSet::new();
+                let mut event_sub = local_broadcast_sender.subscribe();
+                while local_running.load(Ordering::Relaxed) {
+                    tokio::select! {
+                         Some((addr, sub, sender)) = receiver.recv() => {
+                            info!("add or update the subscriber with addr 0x{}", hex::encode(addr.as_ref()));
+                            //TODO limit the max address count
+                            subscribers.insert(addr, (sender, sub));
+                            info!("subscribers len : {}", subscribers.len());
+                        }
+                        Ok(event) = event_sub.recv() => {
+                            info!("receive event {:?}", event);
+                            for (key , (sender, sub)) in subscribers.iter() {
+                                if sender.is_closed() {
+                                    to_be_removed.insert(key.clone());
+                                    warn!("the channel has been closed by client for addr 0x{}", hex::encode(key.as_ref()));
+                                    continue;
+                                }
+                                for idx in 0..sub.topics.len() {
+                                    if sub.topics[idx] != EventTypeV2::Block as i32 {
+                                        continue;
+                                    }
+                                    match sender.try_send(Ok(event.clone())) {
+                                        Ok(_) => {
+                                            info!("send event to addr 0x{}", hex::encode(key.as_ref()));
+                                            break;
+                                        }
+                                        Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => {
+                                            // retry?
+                                            // TODO
+                                            warn!("the channel is full for addr 0x{}", hex::encode(key.as_ref()));
+                                        }
+                                        Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => {
+                                            // remove the address
+                                            to_be_removed.insert(key.clone());
+                                            warn!("the channel has been closed by client for addr 0x{}", hex::encode(key.as_ref()));
+                                        }
+
+                                    }
+                                }
+                            }
+                        },
+                        else => {
+                            info!("unexpected channel update");
+                            // reconnect in 5 seconds
+                            sleep(TokioDuration::from_millis(1000 * 5)).await;
+                            break;
+                        }
+
+                    }
+                    for k in to_be_removed.iter() {
+                        subscribers.remove(k);
+                    }
+                    to_be_removed.clear();
+                }
+            }
+            info!("exit the keep subscription thread");
+        });
+        Ok(())
+    }
 }
 
 #[tonic::async_trait]
 impl StorageNode for StorageNodeV2Impl {
+    type SubscribeStream = ReceiverStream<std::result::Result<EventMessageV2, Status>>;
+    /// add subscription to the light node
+    async fn subscribe(
+        &self,
+        request: Request<SubscribeRequest>,
+    ) -> std::result::Result<Response<Self::SubscribeStream>, Status> {
+        info!("receive subscribe request");
+        let r = request.into_inner();
+        let sender = self.sender.clone();
+        info!("sender is close: {}", sender.is_closed());
+        let account_id = DB3Verifier::verify(r.payload.as_ref(), r.signature.as_ref())
+            .map_err(|e| Status::internal(format!("bad signature for {e}")))?;
+        let payload = SubscriptionV2::decode(r.payload.as_ref()).map_err(|e| {
+            Status::internal(format!("fail to decode open session request for {e} "))
+        })?;
+        info!(
+            "add subscriber for addr 0x{}",
+            hex::encode(account_id.addr.as_ref())
+        );
+        info!("payload {:?}", payload);
+        info!("sender {:?}", sender);
+        let (msg_sender, msg_receiver) =
+            tokio::sync::mpsc::channel::<std::result::Result<EventMessageV2, Status>>(10);
+        sender
+            .try_send((account_id.addr, payload, msg_sender))
+            .map_err(|e| Status::internal(format!("fail to add subscriber for {e}")))?;
+        Ok(Response::new(ReceiverStream::new(msg_receiver)))
+    }
     async fn get_collection_of_database(
         &self,
         request: Request<GetCollectionOfDatabaseRequest>,
diff --git a/src/proto/proto/db3_storage.proto b/src/proto/proto/db3_storage.proto
index 2fdc1119..af47151c 100644
--- a/src/proto/proto/db3_storage.proto
+++ b/src/proto/proto/db3_storage.proto
@@ -21,7 +21,6 @@ package db3_storage_proto;
 import "db3_database_v2.proto";
 import "db3_mutation_v2.proto";
 import "db3_rollup.proto";
-
 message SendMutationRequest {
   // a hex signature string
   string signature = 1;
@@ -54,15 +53,30 @@ message GetNonceResponse {
 }
 
 message SubscribeRequest {
-  string signature = 1;
+  // a hex signature string
+  bytes signature = 1;
+  // the payload of topic
   bytes payload = 2;
 }
-
+// the node will dispatch a block event when a new block has been proposed
 message BlockEvent {
   uint64 block_id = 1;
-  uint64 mutation_count = 2;
+  uint32 mutation_count = 2;
+}
+enum EventType {
+  Block = 0;
+  Mutation = 1;
+  Query = 2;
+}
+message Subscription {
+  repeated EventType topics = 1;
+}
+message EventMessage {
+  EventType type = 1;
+  oneof event {
+    BlockEvent block_event = 3;
+  }
 }
-
 message GetMutationHeaderRequest {
   uint64 block_id = 1;
   uint32 order_id = 2;
@@ -126,4 +140,5 @@ service StorageNode {
   rpc ScanRollupRecord(ScanRollupRecordRequest) returns (ScanRollupRecordResponse) {}
   rpc GetDatabaseOfOwner(GetDatabaseOfOwnerRequest) returns (GetDatabaseOfOwnerResponse) {}
   rpc GetCollectionOfDatabase(GetCollectionOfDatabaseRequest) returns (GetCollectionOfDatabaseResponse) {}
+  rpc Subscribe(SubscribeRequest) returns (stream EventMessage) {}
 }
diff --git a/src/sdk/src/lib.rs b/src/sdk/src/lib.rs
index ba9ef36f..d44fe894 100644
--- a/src/sdk/src/lib.rs
+++ b/src/sdk/src/lib.rs
@@ -20,3 +20,4 @@ pub mod mutation_sdk;
 #[cfg(test)]
 pub mod sdk_test;
 pub mod store_sdk;
+pub mod store_sdk_v2;
diff --git a/src/sdk/src/store_sdk_v2.rs b/src/sdk/src/store_sdk_v2.rs
new file mode 100644
index 00000000..723b0a18
--- /dev/null
+++ b/src/sdk/src/store_sdk_v2.rs
@@ -0,0 +1,119 @@
+//
+// bill_sdk.rs
+// Copyright (C) 2022 db3.network Author imotai <codego.me@gmail.com>
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+
+use bytes::BytesMut;
+use chrono::Utc;
+use db3_crypto::{db3_address::DB3Address, db3_signer::Db3MultiSchemeSigner};
+use db3_proto::db3_storage_proto::event_message::Event as EventV2;
+use db3_proto::db3_storage_proto::{
+    storage_node_client::StorageNodeClient as StorageNodeV2Client, SubscribeRequest,
+};
+use db3_proto::db3_storage_proto::{
+    EventMessage as EventMessageV2, EventType as EventTypeV2, Subscription as SubscriptionV2,
+};
+
+use ethers::core::types::{
+    transaction::eip712::{EIP712Domain, TypedData, Types},
+    Bytes,
+};
+
+use hex;
+use num_traits::cast::FromPrimitive;
+use prost::Message;
+use std::collections::BTreeMap;
+use std::sync::atomic::AtomicBool;
+use std::sync::Arc;
+use tonic::{Status, Streaming};
+
+pub struct StoreSDKV2 {
+    client: Arc<StorageNodeV2Client<tonic::transport::Channel>>,
+    signer: Db3MultiSchemeSigner,
+}
+
+impl StoreSDKV2 {
+    pub fn new(
+        client: Arc<StorageNodeV2Client<tonic::transport::Channel>>,
+        signer: Db3MultiSchemeSigner,
+    ) -> Self {
+        Self { client, signer }
+    }
+
+    pub async fn subscribe_event_message(
+        &mut self,
+    ) -> Result<tonic::Response<Streaming<EventMessageV2>>, Status> {
+        let sub = SubscriptionV2 {
+            topics: vec![EventTypeV2::Block.into()],
+        };
+        let mut buf = BytesMut::with_capacity(1024 * 4);
+        sub.encode(&mut buf)
+            .map_err(|e| Status::internal(format!("Fail to encode subscription {e}")))?;
+        let buf = buf.freeze();
+        let signature = self
+            .signer
+            .sign(buf.as_ref())
+            .map_err(|e| Status::internal(format!("Fail to sign subscription {e}")))?;
+        let req = SubscribeRequest {
+            signature: signature.as_ref().to_vec().to_owned(),
+            payload: buf.as_ref().to_vec().to_owned(),
+        };
+        let mut client = self.client.as_ref().clone();
+        client.subscribe(req).await
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::sdk_test;
+    use std::fmt::format;
+    use std::time::Duration;
+    use tonic::transport::Endpoint;
+
+    async fn subscribe_event_message_flow(
+        use_typed_format: bool,
+        client: Arc<StorageNodeV2Client<tonic::transport::Channel>>,
+        counter: i64,
+    ) {
+        let (_, signer) = sdk_test::gen_secp256k1_signer(counter);
+        let mut sdk = StoreSDKV2::new(client, signer);
+        let res: Result<tonic::Response<Streaming<EventMessageV2>>, Status> =
+            sdk.subscribe_event_message().await;
+        println!("res {:?}", res);
+        assert!(res.is_ok(), "{:?}", res);
+        let two_second = Duration::from_millis(2000);
+        std::thread::sleep(two_second);
+        let mut stream = res.unwrap().into_inner();
+        if let Some(next_message) = stream.message().await.unwrap() {
+            let event_message: EventMessageV2 = next_message;
+            assert_eq!(EventTypeV2::Block as i32, event_message.r#type);
+            assert!(format!("{:?}", event_message)
+                .contains("r#type: Block, event: Some(BlockEvent(BlockEvent { block_id:"));
+        } else {
+            assert!(false);
+        }
+    }
+
+    #[tokio::test]
+    async fn subscribe_event_message_ut() {
+        let ep = "http://127.0.0.1:26619";
+        let rpc_endpoint = Endpoint::new(ep.to_string()).unwrap();
+        let channel = rpc_endpoint.connect_lazy();
+        let client = Arc::new(StorageNodeV2Client::new(channel));
+
+        subscribe_event_message_flow(false, client.clone(), 300).await;
+    }
+}
diff --git a/src/storage/src/mutation_store.rs b/src/storage/src/mutation_store.rs
index 9f9a4e71..fa88d4d2 100644
--- a/src/storage/src/mutation_store.rs
+++ b/src/storage/src/mutation_store.rs
@@ -197,12 +197,13 @@ impl MutationStore {
         Ok(mutation_headers)
     }
 
-    pub fn increase_block(&self) -> Result<()> {
+    pub fn increase_block(&self) -> Result<(u64, u32)> {
         match self.block_state.lock() {
             Ok(mut state) => {
+                let (block, order) = (state.block, state.order);
                 state.block += 1;
                 state.order = 0;
-                Ok(())
+                Ok((block, order))
             }
             Err(e) => Err(DB3Error::WriteStoreError(format!("{e}"))),
         }


Test patch: