feat: RPC for subscriptions

Frando · Frando · commit cb207ad8ee42 · 2024-09-09T12:07:24.000+02:00
diff --git a/iroh-willow/src/engine/actor.rs b/iroh-willow/src/engine/actor.rs
@@ -1,7 +1,7 @@
 use std::{sync::Arc, thread::JoinHandle};
 
 use anyhow::Result;
-use futures_lite::stream::Stream;
+use futures_lite::{stream::Stream, StreamExt};
 use iroh_base::key::NodeId;
 use tokio::{
     sync::{mpsc, oneshot},
@@ -16,13 +16,16 @@ use crate::{
     net::ConnHandle,
     proto::{
         data_model::{AuthorisedEntry, Path, SubspaceId},
-        grouping::Range3d,
+        grouping::{Area, Range3d},
         keys::{NamespaceId, NamespaceKind, UserId, UserSecretKey},
         meadowcap::{self, AccessMode},
     },
     session::{intents::Intent, run_session, Error, EventSender, SessionHandle},
     store::{
-        traits::{EntryOrigin, EntryReader, EntryStorage, SecretStorage, Storage},
+        traits::{
+            EntryOrigin, EntryReader, EntryStorage, SecretStorage, Storage, StoreEvent,
+            SubscribeParams,
+        },
         Store,
     },
 };
@@ -211,6 +214,42 @@ impl ActorHandle {
         reply_rx.await?;
         Ok(())
     }
+
+    pub async fn subscribe_area(
+        &self,
+        namespace: NamespaceId,
+        area: Area,
+        params: SubscribeParams,
+        sender: mpsc::Sender<StoreEvent>,
+    ) -> Result<()> {
+        self.send(Input::SubscribeArea {
+            namespace,
+            area,
+            params,
+            sender,
+        })
+        .await?;
+        Ok(())
+    }
+
+    pub async fn resume_subscription(
+        &self,
+        progress_id: u64,
+        namespace: NamespaceId,
+        area: Area,
+        params: SubscribeParams,
+        sender: mpsc::Sender<StoreEvent>,
+    ) -> Result<()> {
+        self.send(Input::ResumeSubscription {
+            progress_id,
+            namespace,
+            area,
+            params,
+            sender,
+        })
+        .await?;
+        Ok(())
+    }
 }
 
 impl Drop for ActorHandle {
@@ -308,6 +347,19 @@ pub enum Input {
         #[debug(skip)]
         reply: Option<oneshot::Sender<()>>,
     },
+    SubscribeArea {
+        namespace: NamespaceId,
+        area: Area,
+        params: SubscribeParams,
+        sender: mpsc::Sender<StoreEvent>,
+    },
+    ResumeSubscription {
+        progress_id: u64,
+        namespace: NamespaceId,
+        area: Area,
+        params: SubscribeParams,
+        sender: mpsc::Sender<StoreEvent>,
+    },
 }
 
 #[derive(Debug)]
@@ -488,6 +540,44 @@ impl<S: Storage> Actor<S> {
                 let res = self.store.auth().resolve_interests(interests);
                 send_reply(reply, res.map_err(anyhow::Error::from))
             }
+            Input::SubscribeArea {
+                namespace,
+                area,
+                params,
+                sender,
+            } => {
+                let store = self.store.clone();
+                self.tasks.spawn_local(async move {
+                    let mut stream = store.entries().subscribe_area(namespace, area, params);
+                    while let Some(event) = stream.next().await {
+                        if let Err(_) = sender.send(event).await {
+                            break;
+                        }
+                    }
+                });
+                Ok(())
+            }
+            Input::ResumeSubscription {
+                progress_id,
+                namespace,
+                area,
+                params,
+                sender,
+            } => {
+                let store = self.store.clone();
+                self.tasks.spawn_local(async move {
+                    let mut stream =
+                        store
+                            .entries()
+                            .resume_subscription(progress_id, namespace, area, params);
+                    while let Some(event) = stream.next().await {
+                        if let Err(_) = sender.send(event).await {
+                            break;
+                        }
+                    }
+                });
+                Ok(())
+            }
         }
     }
 }
diff --git a/iroh-willow/src/store/memory.rs b/iroh-willow/src/store/memory.rs
@@ -257,16 +257,12 @@ impl EntryStore {
             }
         }
         for i in to_prune {
-            let pruned = entries.remove(i).into_parts().0;
-            self.events.insert(|id| {
+            let pruned = entries.remove(i);
+            self.events.insert(move |id| {
                 StoreEvent::Pruned(
                     id,
                     traits::PruneEvent {
-                        pruned: (
-                            pruned.namespace_id().clone(),
-                            pruned.subspace_id().clone(),
-                            pruned.path().clone(),
-                        ),
+                        pruned,
                         by: entry.clone(),
                     },
                 )
@@ -404,7 +400,7 @@ impl<T> Default for EventQueue<T> {
 }
 
 impl<T: Clone> EventQueue<T> {
-    fn insert(&mut self, f: impl Fn(u64) -> T) {
+    fn insert(&mut self, f: impl FnOnce(u64) -> T) {
         let progress_id = self.next_progress_id();
         let event = f(progress_id);
         self.events.push_back(event);
diff --git a/iroh-willow/src/store/traits.rs b/iroh-willow/src/store/traits.rs
@@ -4,11 +4,14 @@ use std::fmt::Debug;
 
 use anyhow::Result;
 use futures_lite::Stream;
+use serde::{Deserialize, Serialize};
 
 use crate::{
     interest::{CapSelector, CapabilityPack},
     proto::{
-        data_model::{AuthorisedEntry, Entry, NamespaceId, Path, SubspaceId, WriteCapability},
+        data_model::{
+            self, AuthorisedEntry, Entry, NamespaceId, Path, SubspaceId, WriteCapability,
+        },
         grouping::{Area, Range3d},
         keys::{NamespaceSecretKey, NamespaceSignature, UserId, UserSecretKey, UserSignature},
         meadowcap::{self, ReadAuthorisation},
@@ -209,10 +212,14 @@ pub trait CapsStorage: Debug + Clone {
 
 /// An event which took place within a [`EntryStorage`].
 /// Each event includes a *progress ID* which can be used to *resume* a subscription at any point in the future.
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
 pub enum StoreEvent {
     /// A new entry was ingested.
-    Ingested(u64, AuthorisedEntry, EntryOrigin),
+    Ingested(
+        u64,
+        #[serde(with = "data_model::serde_encoding::authorised_entry")] AuthorisedEntry,
+        EntryOrigin,
+    ),
     // PayloadForgotten(u64, PD),
     /// An entry was pruned via prefix pruning.
     Pruned(u64, PruneEvent),
@@ -249,9 +256,8 @@ impl StoreEvent {
             }
             StoreEvent::Pruned(_, PruneEvent { pruned, by: _ }) => {
                 if !params.ingest_only
-                    && pruned.0 == namespace_id
-                    && area.subspace().includes(&pruned.1)
-                    && area.path().is_prefix_of(&pruned.2)
+                    && *pruned.entry().namespace_id() == namespace_id
+                    && area.includes_entry(pruned.entry())
                 {
                     true
                 } else {
@@ -263,16 +269,17 @@ impl StoreEvent {
 }
 
 /// Describes an [`AuthorisedEntry`] which was pruned and the [`AuthorisedEntry`] which triggered the pruning.
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct PruneEvent {
-    /// The subspace ID and path of the entry which was pruned.
-    pub pruned: (NamespaceId, SubspaceId, Path),
+    #[serde(with = "data_model::serde_encoding::authorised_entry")]
+    pub pruned: AuthorisedEntry,
     /// The entry which triggered the pruning.
+    #[serde(with = "data_model::serde_encoding::authorised_entry")]
     pub by: AuthorisedEntry,
 }
 
 /// The origin of an entry ingestion event.
-#[derive(Debug, Clone, Copy, Eq, PartialEq)]
+#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
 pub enum EntryOrigin {
     /// The entry was probably created on this machine.
     Local,
@@ -287,7 +294,7 @@ pub enum EntryChannel {
 }
 
 /// Describes which entries to ignore during a query.
-#[derive(Debug, Default)]
+#[derive(Debug, Default, Serialize, Deserialize)]
 pub struct SubscribeParams {
     /// Omit entries whose payload is the empty string.
     pub ignore_empty_payloads: bool,
diff --git a/iroh/src/client/spaces.rs b/iroh/src/client/spaces.rs
@@ -38,6 +38,7 @@ use iroh_willow::{
         intents::{serde_encoding::Event, Completion, IntentUpdate},
         SessionInit,
     },
+    store::traits::{StoreEvent, SubscribeParams},
 };
 use ref_cast::RefCast;
 use serde::{Deserialize, Serialize};
@@ -353,14 +354,39 @@ impl Space {
         })
     }
 
-    /// TODO
-    pub fn subscribe(&self, _area: Area) {
-        todo!()
+    /// Subscribe to events concerning entries included by an `Area`.
+    pub async fn subscribe_area(
+        &self,
+        area: Area,
+        params: SubscribeParams,
+    ) -> Result<impl Stream<Item = Result<StoreEvent>>> {
+        let req = SubscribeRequest {
+            namespace: self.namespace_id,
+            area,
+            params,
+            initial_progress_id: None,
+        };
+        let stream = self.rpc.try_server_streaming(req).await?;
+        let stream = stream.map(|item| item.map_err(anyhow::Error::from));
+        Ok(stream)
     }
 
-    /// TODO
-    pub fn subscribe_offset(&self, _area: Area, _offset: u64) {
-        todo!()
+    /// Resume a subscription using a progress ID obtained from a previous subscription.
+    pub async fn resume_subscription(
+        &self,
+        progress_id: u64,
+        area: Area,
+        params: SubscribeParams,
+    ) -> Result<impl Stream<Item = Result<StoreEvent>>> {
+        let req = SubscribeRequest {
+            namespace: self.namespace_id,
+            area,
+            params,
+            initial_progress_id: Some(progress_id),
+        };
+        let stream = self.rpc.try_server_streaming(req).await?;
+        let stream = stream.map(|item| item.map_err(anyhow::Error::from));
+        Ok(stream)
     }
 }
 
diff --git a/iroh/src/node/rpc/spaces.rs b/iroh/src/node/rpc/spaces.rs
@@ -155,6 +155,32 @@ impl<D: Store> Handler<D> {
                 .await
             }
             SyncWithPeerUpdate(_) => Err(RpcServerError::UnexpectedStartMessage),
+            Subscribe(msg) => {
+                chan.try_server_streaming(msg, self, |handler, req| async move {
+                    let (tx, rx) = mpsc::channel(1024);
+                    if let Some(progress_id) = req.initial_progress_id {
+                        handler
+                            .spaces()?
+                            .resume_subscription(
+                                progress_id,
+                                req.namespace,
+                                req.area,
+                                req.params,
+                                tx,
+                            )
+                            .await
+                            .map_err(map_err)?;
+                    } else {
+                        handler
+                            .spaces()?
+                            .subscribe_area(req.namespace, req.area, req.params, tx)
+                            .await
+                            .map_err(map_err)?;
+                    }
+                    Ok(ReceiverStream::new(rx).map(Ok))
+                })
+                .await
+            }
         }
     }
 }
diff --git a/iroh/src/rpc_protocol/spaces.rs b/iroh/src/rpc_protocol/spaces.rs
@@ -9,14 +9,15 @@ use iroh_willow::{
             self, serde_encoding::SerdeAuthorisedEntry, AuthorisedEntry, Entry, NamespaceId, Path,
             SubspaceId,
         },
-        grouping::{self, Range3d},
+        grouping::{self, Area, Range3d},
         keys::{NamespaceKind, UserId},
         meadowcap::{self, AccessMode, SecretKey},
     },
     session::{
         intents::{serde_encoding::Event, IntentUpdate},
         SessionInit,
     },
+    store::traits::{StoreEvent, SubscribeParams},
 };
 use nested_enum_utils::enum_conversions;
 use quic_rpc_derive::rpc_requests;
@@ -47,11 +48,11 @@ pub enum Request {
     DelegateCaps(DelegateCapsRequest),
     #[rpc(response = RpcResult<ImportCapsResponse>)]
     ImportCaps(ImportCapsRequest),
-    // #[rpc(response = RpcResult<ResolveInterestsResponse>)]
-    // ResolveInterests(ResolveInterestsRequest),
     #[bidi_streaming(update = SyncWithPeerUpdate, response = RpcResult<SyncWithPeerResponse>)]
     SyncWithPeer(SyncWithPeerRequest),
     SyncWithPeerUpdate(SyncWithPeerUpdate),
+    #[try_server_streaming(create_error = RpcError, item_error = RpcError, item = StoreEvent)]
+    Subscribe(SubscribeRequest),
 }
 
 #[allow(missing_docs)]
@@ -67,8 +68,8 @@ pub enum Response {
     CreateUser(RpcResult<CreateUserResponse>),
     DelegateCaps(RpcResult<DelegateCapsResponse>),
     ImportCaps(RpcResult<ImportCapsResponse>),
-    // ResolveInterests(RpcResult<ResolveInterestsResponse>),
     SyncWithPeer(RpcResult<SyncWithPeerResponse>),
+    Subscribe(RpcResult<StoreEvent>),
 }
 
 #[derive(Debug, Serialize, Deserialize)]
@@ -139,9 +140,7 @@ pub struct GetEntryRequest {
 }
 
 #[derive(Debug, Serialize, Deserialize)]
-pub struct GetEntryResponse(
-    pub Option<SerdeAuthorisedEntry>, // #[serde(with = "data_model::serde_encoding::authorised_entry")] pub AuthorisedEntry,
-);
+pub struct GetEntryResponse(pub Option<SerdeAuthorisedEntry>);
 
 #[derive(Debug, Serialize, Deserialize)]
 pub struct CreateNamespaceRequest {
@@ -192,6 +191,15 @@ pub enum SyncWithPeerResponse {
     Event(Event),
 }
 
+#[derive(Debug, Serialize, Deserialize)]
+pub struct SubscribeRequest {
+    pub namespace: NamespaceId,
+    #[serde(with = "grouping::serde_encoding::area")]
+    pub area: Area,
+    pub params: SubscribeParams,
+    pub initial_progress_id: Option<u64>,
+}
+
 /// Either a complete [`Entry`] or a [`FullEntryForm`].
 #[derive(Debug, Serialize, Deserialize)]
 pub enum EntryOrForm {
