RPC subscriptions for new slot notifications (#7114)

* feat: slot notifications via pubsub rpc w/ tests

core/src/rpc_subscriptions.rs

@@ -15,6 +15,13 @@ use std::sync::{Arc, RwLock};
 
 pub type Confirmations = usize;
 
+#[derive(Serialize, Clone)]
+pub struct SlotInfo {
+    pub slot: Slot,
+    pub parent: Slot,
+    pub root: Slot,
+}
+
 type RpcAccountSubscriptions =
     RwLock<HashMap<Pubkey, HashMap<SubscriptionId, (Sink<Account>, Confirmations)>>>;
 type RpcProgramSubscriptions =
@@ -22,6 +29,7 @@ type RpcProgramSubscriptions =
 type RpcSignatureSubscriptions = RwLock<
     HashMap<Signature, HashMap<SubscriptionId, (Sink<transaction::Result<()>>, Confirmations)>>,
 >;
+type RpcSlotSubscriptions = RwLock<HashMap<SubscriptionId, Sink<SlotInfo>>>;
 
 fn add_subscription<K, S>(
     subscriptions: &mut HashMap<K, HashMap<SubscriptionId, (Sink<S>, Confirmations)>>,
@@ -119,7 +127,7 @@ fn check_confirmations_and_notify<K, S, F, N, X>(
     }
 }
 
-fn notify_account<S>(result: Option<(S, u64)>, sink: &Sink<S>, root: u64)
+fn notify_account<S>(result: Option<(S, Slot)>, sink: &Sink<S>, root: Slot)
 where
     S: Clone + Serialize,
 {
@@ -130,7 +138,7 @@ where
     }
 }
 
-fn notify_signature<S>(result: Option<S>, sink: &Sink<S>, _root: u64)
+fn notify_signature<S>(result: Option<S>, sink: &Sink<S>, _root: Slot)
 where
     S: Clone + Serialize,
 {
@@ -139,7 +147,7 @@ where
     }
 }
 
-fn notify_program(accounts: Vec<(Pubkey, Account)>, sink: &Sink<(String, Account)>, _root: u64) {
+fn notify_program(accounts: Vec<(Pubkey, Account)>, sink: &Sink<(String, Account)>, _root: Slot) {
     for (pubkey, account) in accounts.iter() {
         sink.notify(Ok((pubkey.to_string(), account.clone())))
             .wait()
@@ -151,6 +159,7 @@ pub struct RpcSubscriptions {
     account_subscriptions: RpcAccountSubscriptions,
     program_subscriptions: RpcProgramSubscriptions,
     signature_subscriptions: RpcSignatureSubscriptions,
+    slot_subscriptions: RpcSlotSubscriptions,
 }
 
 impl Default for RpcSubscriptions {
@@ -159,6 +168,7 @@ impl Default for RpcSubscriptions {
             account_subscriptions: RpcAccountSubscriptions::default(),
             program_subscriptions: RpcProgramSubscriptions::default(),
             signature_subscriptions: RpcSignatureSubscriptions::default(),
+            slot_subscriptions: RpcSlotSubscriptions::default(),
         }
     }
 }
@@ -291,6 +301,26 @@ impl RpcSubscriptions {
             self.check_signature(signature, current_slot, bank_forks);
         }
     }
+
+    pub fn add_slot_subscription(&self, sub_id: &SubscriptionId, sink: &Sink<SlotInfo>) {
+        let mut subscriptions = self.slot_subscriptions.write().unwrap();
+        subscriptions.insert(sub_id.clone(), sink.clone());
+    }
+
+    pub fn remove_slot_subscription(&self, id: &SubscriptionId) -> bool {
+        let mut subscriptions = self.slot_subscriptions.write().unwrap();
+        subscriptions.remove(id).is_some()
+    }
+
+    pub fn notify_slot(&self, slot: Slot, parent: Slot, root: Slot) {
+        info!("notify_slot!! {} from {} (root={})", slot, parent, root);
+        let subscriptions = self.slot_subscriptions.read().unwrap();
+        for (_, sink) in subscriptions.iter() {
+            sink.notify(Ok(SlotInfo { slot, parent, root }))
+                .wait()
+                .unwrap();
+        }
+    }
 }
 
 #[cfg(test)]
@@ -463,4 +493,40 @@ mod tests {
             .unwrap()
             .contains_key(&signature));
     }
+    #[test]
+    fn test_check_slot_subscribe() {
+        let (subscriber, _id_receiver, mut transport_receiver) =
+            Subscriber::new_test("slotNotification");
+        let sub_id = SubscriptionId::Number(0 as u64);
+        let sink = subscriber.assign_id(sub_id.clone()).unwrap();
+        let subscriptions = RpcSubscriptions::default();
+        subscriptions.add_slot_subscription(&sub_id, &sink);
+
+        assert!(subscriptions
+            .slot_subscriptions
+            .read()
+            .unwrap()
+            .contains_key(&sub_id));
+
+        subscriptions.notify_slot(0, 0, 0);
+        let string = transport_receiver.poll();
+        if let Async::Ready(Some(response)) = string.unwrap() {
+            let expected_res = SlotInfo {
+                parent: 0,
+                slot: 0,
+                root: 0,
+            };
+            let expected_res_str =
+                serde_json::to_string(&serde_json::to_value(expected_res).unwrap()).unwrap();
+            let expected = format!(r#"{{"jsonrpc":"2.0","method":"slotNotification","params":{{"result":{},"subscription":0}}}}"#, expected_res_str);
+            assert_eq!(expected, response);
+        }
+
+        subscriptions.remove_slot_subscription(&sub_id);
+        assert!(!subscriptions
+            .slot_subscriptions
+            .read()
+            .unwrap()
+            .contains_key(&sub_id));
+    }
 }