implement variant of subscription that returns finalized storage changes (#237)

Signed-off-by: Gregory Hill <gregorydhill@outlook.com>

src/subscription.rs

@@ -16,7 +16,14 @@
 
 use jsonrpsee_types::error::Error as RpcError;
 use jsonrpsee_ws_client::WsSubscription as Subscription;
-use sp_core::storage::StorageChangeSet;
+use sp_core::{
+    storage::{
+        StorageChangeSet,
+        StorageKey,
+    },
+    twox_128,
+};
+use sp_runtime::traits::Header;
 use std::collections::VecDeque;
 
 use crate::{
@@ -30,13 +37,14 @@ use crate::{
         system::Phase,
         Event,
     },
+    rpc::Rpc,
     runtimes::Runtime,
 };
 
 /// Event subscription simplifies filtering a storage change set stream for
 /// events of interest.
 pub struct EventSubscription<'a, T: Runtime> {
-    subscription: Subscription<StorageChangeSet<T::Hash>>,
+    subscription: EventStorageSubscription<T>,
     decoder: &'a EventsDecoder<T>,
     block: Option<T::Hash>,
     extrinsic: Option<usize>,
@@ -48,7 +56,7 @@ pub struct EventSubscription<'a, T: Runtime> {
 impl<'a, T: Runtime> EventSubscription<'a, T> {
     /// Creates a new event subscription.
     pub fn new(
-        subscription: Subscription<StorageChangeSet<T::Hash>>,
+        subscription: EventStorageSubscription<T>,
         decoder: &'a EventsDecoder<T>,
     ) -> Self {
         Self {
@@ -132,3 +140,75 @@ impl<'a, T: Runtime> EventSubscription<'a, T> {
         }
     }
 }
+
+pub(crate) struct SystemEvents(StorageKey);
+
+impl SystemEvents {
+    pub(crate) fn new() -> Self {
+        let mut storage_key = twox_128(b"System").to_vec();
+        storage_key.extend(twox_128(b"Events").to_vec());
+        log::debug!("Events storage key {:?}", hex::encode(&storage_key));
+        Self(StorageKey(storage_key))
+    }
+}
+
+impl From<SystemEvents> for StorageKey {
+    fn from(key: SystemEvents) -> Self {
+        key.0
+    }
+}
+
+/// Event subscription to only fetch finalized storage changes.
+pub struct FinalizedEventStorageSubscription<T: Runtime> {
+    rpc: Rpc<T>,
+    subscription: Subscription<T::Header>,
+    storage_changes: VecDeque<StorageChangeSet<T::Hash>>,
+    storage_key: StorageKey,
+}
+
+impl<T: Runtime> FinalizedEventStorageSubscription<T> {
+    /// Creates a new finalized event storage subscription.
+    pub fn new(rpc: Rpc<T>, subscription: Subscription<T::Header>) -> Self {
+        Self {
+            rpc,
+            subscription,
+            storage_changes: Default::default(),
+            storage_key: SystemEvents::new().into(),
+        }
+    }
+
+    /// Gets the next change_set.
+    pub async fn next(&mut self) -> Option<StorageChangeSet<T::Hash>> {
+        loop {
+            if let Some(storage_change) = self.storage_changes.pop_front() {
+                return Some(storage_change)
+            }
+            let header: T::Header = self.subscription.next().await?;
+            if let Ok(storage_changes) = self
+                .rpc
+                .query_storage_at(&[self.storage_key.clone()], Some(header.hash()))
+                .await
+            {
+                self.storage_changes.extend(storage_changes);
+            }
+        }
+    }
+}
+
+/// Wrapper over imported and finalized event subscriptions.
+pub enum EventStorageSubscription<T: Runtime> {
+    /// Events that are InBlock
+    Imported(Subscription<StorageChangeSet<T::Hash>>),
+    /// Events that are Finalized
+    Finalized(FinalizedEventStorageSubscription<T>),
+}
+
+impl<T: Runtime> EventStorageSubscription<T> {
+    /// Gets the next change_set from the subscription.
+    pub async fn next(&mut self) -> Option<StorageChangeSet<T::Hash>> {
+        match self {
+            Self::Imported(event_sub) => event_sub.next().await,
+            Self::Finalized(event_sub) => event_sub.next().await,
+        }
+    }
+}