feat: add rpc to subscribe store changes

Author: doitian

crates/fiber-bin/src/main.rs

@@ -354,6 +354,7 @@ pub async fn main() -> Result<(), ExitMessage> {
                 cch_actor,
                 store,
                 network_graph,
+                root_actor.get_cell(),
                 #[cfg(debug_assertions)] ckb_chain_actor,
                 #[cfg(debug_assertions)] rpc_dev_module_commitment_txs,
             )

crates/fiber-lib/src/rpc/biscuit.rs

@@ -78,6 +78,7 @@ fn build_rules() -> HashMap<&'static str, AuthRule> {
     b.rule("send_btc", r#"allow if write("cch");"#);
     b.rule("receive_btc", r#"allow if read("cch");"#);
     b.rule("get_cch_order", r#"allow if read("cch");"#);
+    b.rule("subscribe_store_changes", r#"allow if read("cch");"#);
     // channels
     b.rule("open_channel", r#"allow if write("channels");"#);
     b.rule("accept_channel", r#"allow if write("channels");"#);

crates/fiber-lib/src/rpc/config.rs

@@ -1,9 +1,10 @@
 use clap_serde_derive::ClapSerde;
 
 #[cfg(not(feature = "watchtower"))]
-const DEFAULT_ENABLED_MODULES: &str = "cch,channel,graph,payment,info,invoice,peer";
+const DEFAULT_ENABLED_MODULES: &str = "cch,channel,graph,payment,info,invoice,peer,pubsub";
 #[cfg(feature = "watchtower")]
-const DEFAULT_ENABLED_MODULES: &str = "cch,channel,graph,payment,info,invoice,peer,watchtower";
+const DEFAULT_ENABLED_MODULES: &str =
+    "cch,channel,graph,payment,info,invoice,peer,pubsub,watchtower";
 
 #[derive(ClapSerde, Debug, Clone)]
 pub struct RpcConfig {

crates/fiber-lib/src/rpc/mod.rs

@@ -35,6 +35,7 @@ pub mod server {
     use crate::rpc::payment::PaymentRpcServer;
     use crate::rpc::payment::PaymentRpcServerImpl;
     use crate::rpc::peer::{PeerRpcServer, PeerRpcServerImpl};
+    use crate::store::pub_sub::{register_pub_sub_rpc, Subscribe};
     use crate::{
         cch::CchMessage,
         fiber::{
@@ -60,7 +61,7 @@ pub mod server {
     };
     use jsonrpsee::ws_client::RpcServiceBuilder;
     use jsonrpsee::{Methods, RpcModule};
-    use ractor::ActorRef;
+    use ractor::{ActorCell, ActorRef};
     #[cfg(debug_assertions)]
     use std::collections::HashMap;
     use std::net::{IpAddr, SocketAddr, ToSocketAddrs};
@@ -80,6 +81,7 @@ pub mod server {
         + GossipMessageStore
         + WatchtowerStore
         + PreimageStore
+        + Subscribe
     {
     }
     #[cfg(feature = "watchtower")]
@@ -90,16 +92,25 @@ pub mod server {
             + GossipMessageStore
             + WatchtowerStore
             + PreimageStore
+            + Subscribe
     {
     }
     #[cfg(not(feature = "watchtower"))]
     pub trait RpcServerStore:
-        ChannelActorStateStore + InvoiceStore + NetworkGraphStateStore + GossipMessageStore
+        ChannelActorStateStore
+        + InvoiceStore
+        + NetworkGraphStateStore
+        + GossipMessageStore
+        + Subscribe
     {
     }
     #[cfg(not(feature = "watchtower"))]
     impl<T> RpcServerStore for T where
-        T: ChannelActorStateStore + InvoiceStore + NetworkGraphStateStore + GossipMessageStore
+        T: ChannelActorStateStore
+            + InvoiceStore
+            + NetworkGraphStateStore
+            + GossipMessageStore
+            + Subscribe
     {
     }
 
@@ -222,6 +233,7 @@ pub mod server {
         cch_actor: Option<ActorRef<CchMessage>>,
         store: S,
         network_graph: Arc<RwLock<NetworkGraph<S>>>,
+        supervisor: ActorCell,
         #[cfg(debug_assertions)] ckb_chain_actor: Option<ActorRef<CkbChainMessage>>,
         #[cfg(debug_assertions)] rpc_dev_module_commitment_txs: Option<
             Arc<RwLock<HashMap<(Hash256, u64), TransactionView>>>,
@@ -258,6 +270,9 @@ pub mod server {
                 .merge(GraphRpcServerImpl::new(network_graph, store.clone()).into_rpc())
                 .unwrap();
         }
+        if config.is_module_enabled("pubsub") {
+            register_pub_sub_rpc(&mut modules, &store, supervisor).await?;
+        }
         if let Some(network_actor) = network_actor {
             if config.is_module_enabled("info") {
                 modules

crates/fiber-lib/src/store/pub_sub/mod.rs

@@ -1,7 +1,9 @@
+mod pub_sub_rpc;
 mod store_with_pub_sub;
 mod subscription;
 
-pub use store_with_pub_sub::StoreWithPubSub;
+pub use pub_sub_rpc::register_pub_sub_rpc;
+pub use store_with_pub_sub::{StoreWithPubSub, Subscribe};
 pub use subscription::{
     InvoiceUpdatedEvent, InvoiceUpdatedPayload, PaymentUpdatedEvent, PaymentUpdatedPayload,
     StorePublisher, StoreUpdatedEvent,

crates/fiber-lib/src/store/pub_sub/pub_sub_rpc.rs

@@ -0,0 +1,90 @@
+use crate::store::pub_sub::Subscribe;
+
+/// Store Pub/Sub via RPC
+use super::StoreUpdatedEvent;
+
+use jsonrpsee::{RpcModule, SubscriptionSink};
+use ractor::{Actor, ActorCell, ActorProcessingErr, ActorRef};
+
+pub struct PubSubServerActor;
+
+#[derive(Default)]
+pub struct PubSubServerState {
+    sinks: Vec<SubscriptionSink>,
+}
+
+pub enum PubSubServerMessage {
+    Publish(StoreUpdatedEvent),
+    AddSink(SubscriptionSink),
+}
+
+impl From<StoreUpdatedEvent> for PubSubServerMessage {
+    fn from(event: StoreUpdatedEvent) -> Self {
+        PubSubServerMessage::Publish(event)
+    }
+}
+
+#[async_trait::async_trait]
+impl Actor for PubSubServerActor {
+    type State = PubSubServerState;
+    type Msg = PubSubServerMessage;
+    type Arguments = ();
+
+    async fn pre_start(
+        &self,
+        _myself: ActorRef<Self::Msg>,
+        _args: Self::Arguments,
+    ) -> Result<Self::State, ActorProcessingErr> {
+        Ok(PubSubServerState::default())
+    }
+
+    async fn handle(
+        &self,
+        _myself: ActorRef<Self::Msg>,
+        message: Self::Msg,
+        state: &mut Self::State,
+    ) -> Result<(), ActorProcessingErr> {
+        match message {
+            PubSubServerMessage::AddSink(sink) => state.sinks.push(sink),
+            PubSubServerMessage::Publish(event) => {
+                let subscription_message =
+                    serde_json::value::to_raw_value(&event).expect("serialize to JSON");
+                let sinks = std::mem::take(&mut state.sinks);
+                for sink in sinks {
+                    if sink.send(subscription_message.clone()).await.is_ok() {
+                        state.sinks.push(sink);
+                    }
+                }
+            }
+        }
+        Ok(())
+    }
+}
+
+const SUBSCRIBE_STORE_CHANGES_NAME: &str = "subscribe_store_changes";
+const SUBSCRIBE_STORE_CHANGES_NOTIF_NAME: &str = "store_changes";
+const UNSUBSCRIBE_STORE_CHANGES_NAME: &str = "unsubscribe_store_changes";
+
+pub async fn register_pub_sub_rpc<S: Subscribe>(
+    modules: &mut RpcModule<()>,
+    publisher: &S,
+    supervisor: ActorCell,
+) -> anyhow::Result<()> {
+    let (pub_sub_actor, _) =
+        ractor::Actor::spawn_linked(None, PubSubServerActor, (), supervisor).await?;
+    publisher.subscribe(Box::new(pub_sub_actor.clone()));
+    modules.register_subscription(
+        SUBSCRIBE_STORE_CHANGES_NAME,
+        SUBSCRIBE_STORE_CHANGES_NOTIF_NAME,
+        UNSUBSCRIBE_STORE_CHANGES_NAME,
+        move |_, pending, _, _| {
+            let pub_sub_actor = pub_sub_actor.clone();
+            async move {
+                let sink = pending.accept().await?;
+                let _ = pub_sub_actor.send_message(PubSubServerMessage::AddSink(sink));
+                Ok(())
+            }
+        },
+    )?;
+    Ok(())
+}

crates/fiber-lib/src/store/pub_sub/store_with_pub_sub.rs

@@ -47,6 +47,16 @@ impl<S> StoreWithPubSub<S> {
     }
 }
 
+pub trait Subscribe {
+    fn subscribe(&self, subscriber: OutputPortSubscriber<StoreUpdatedEvent>);
+}
+
+impl<S> Subscribe for StoreWithPubSub<S> {
+    fn subscribe(&self, subscriber: OutputPortSubscriber<StoreUpdatedEvent>) {
+        self.subscribe(subscriber);
+    }
+}
+
 impl<T: NetworkActorStateStore> NetworkActorStateStoreDeref for StoreWithPubSub<T> {
     type Target = T;
 

crates/fiber-lib/src/store/pub_sub/subscription.rs

@@ -3,13 +3,14 @@
 use std::sync::Arc;
 
 use ractor::{port::OutputPortSubscriber, OutputPort};
+use serde::{Deserialize, Serialize};
 
 use crate::{
     fiber::{payment::PaymentStatus, types::Hash256},
     invoice::CkbInvoiceStatus,
 };
 
-#[derive(Clone, Debug, Eq, PartialEq)]
+#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
 pub enum InvoiceUpdatedPayload {
     /// The invoice is open and can be paid.
     Open,
@@ -28,7 +29,7 @@ pub enum InvoiceUpdatedPayload {
     /// The invoice is paid.
     Paid,
 }
-#[derive(Clone, Debug, Eq, PartialEq)]
+#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
 pub struct InvoiceUpdatedEvent {
     pub invoice_hash: Hash256,
     pub payload: InvoiceUpdatedPayload,
@@ -61,7 +62,7 @@ impl InvoiceUpdatedEvent {
 }
 
 // but with additional information for downstream services.
-#[derive(Clone, Debug, Eq, PartialEq)]
+#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
 pub enum PaymentUpdatedPayload {
     /// initial status, payment session is created, no HTLC is sent
     Created,
@@ -72,7 +73,7 @@ pub enum PaymentUpdatedPayload {
     /// related HTLC is failed
     Failed,
 }
-#[derive(Clone, Debug, Eq, PartialEq)]
+#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
 pub struct PaymentUpdatedEvent {
     pub payment_hash: Hash256,
     pub payload: PaymentUpdatedPayload,
@@ -104,7 +105,7 @@ impl PaymentUpdatedEvent {
 }
 
 /// Message sent from Store to publisher.
-#[derive(Clone, Debug, Eq, PartialEq)]
+#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
 pub enum StoreUpdatedEvent {
     InvoiceUpdated(InvoiceUpdatedEvent),
     PaymentUpdated(PaymentUpdatedEvent),

crates/fiber-lib/src/tests/test_utils.rs

@@ -1563,6 +1563,7 @@ impl NetworkNode {
                     None,
                     store.clone(),
                     network_graph.clone(),
+                    root_actor.get_cell(),
                     #[cfg(debug_assertions)]
                     None,
                     #[cfg(debug_assertions)]

tests/nodes/deployer/config.yml

@@ -20,6 +20,7 @@ rpc:
     - invoice
     - peer
     - watchtower
+    - pubsub
     - dev
 
 cch: