Move EntryStream into its own Tvu stage

src/entry_stream.rs

@@ -3,12 +3,14 @@
 //! real-time access to entries.
 
 use crate::entry::Entry;
+use crate::leader_scheduler::LeaderScheduler;
 use crate::result::Result;
-use chrono::Utc;
+use chrono::{SecondsFormat, Utc};
 use std::io::prelude::*;
 use std::net::Shutdown;
 use std::os::unix::net::UnixStream;
 use std::path::Path;
+use std::sync::{Arc, RwLock};
 
 pub trait EntryStreamHandler {
     fn stream_entries(&mut self, entries: &[Entry]) -> Result<()>;
@@ -16,11 +18,15 @@ pub trait EntryStreamHandler {
 
 pub struct EntryStream {
     pub socket: String,
+    leader_scheduler: Arc<RwLock<LeaderScheduler>>,
 }
 
 impl EntryStream {
-    pub fn new(socket: String) -> Self {
-        EntryStream { socket }
+    pub fn new(socket: String, leader_scheduler: Arc<RwLock<LeaderScheduler>>) -> Self {
+        EntryStream {
+            socket,
+            leader_scheduler,
+        }
     }
 }
 
@@ -29,7 +35,18 @@ impl EntryStreamHandler for EntryStream {
         let mut socket = UnixStream::connect(Path::new(&self.socket))?;
         for entry in entries {
             let json = serde_json::to_string(&entry)?;
-            let payload = format!(r#"{{"dt":"{}","entry":{}}}"#, Utc::now().to_rfc3339(), json);
+            let (slot, slot_leader) = {
+                let leader_scheduler = self.leader_scheduler.read().unwrap();
+                let slot = leader_scheduler.tick_height_to_slot(entry.tick_height);
+                (slot, leader_scheduler.get_leader_for_slot(slot))
+            };
+            let payload = format!(
+                r#"{{"dt":"{}","t":"entry","s":{},"leader_id":"{:?}","entry":{}}}"#,
+                Utc::now().to_rfc3339_opts(SecondsFormat::Nanos, true),
+                slot,
+                slot_leader,
+                json
+            );
             socket.write_all(payload.as_bytes())?;
         }
         socket.shutdown(Shutdown::Write)?;
@@ -39,12 +56,16 @@ impl EntryStreamHandler for EntryStream {
 
 pub struct MockEntryStream {
     pub socket: Vec<String>,
+    leader_scheduler: Arc<RwLock<LeaderScheduler>>,
 }
 
 impl MockEntryStream {
     #[allow(clippy::needless_pass_by_value)]
-    pub fn new(_socket: String) -> Self {
-        MockEntryStream { socket: Vec::new() }
+    pub fn new(_socket: String, leader_scheduler: Arc<RwLock<LeaderScheduler>>) -> Self {
+        MockEntryStream {
+            socket: Vec::new(),
+            leader_scheduler,
+        }
     }
 }
 
@@ -52,7 +73,18 @@ impl EntryStreamHandler for MockEntryStream {
     fn stream_entries(&mut self, entries: &[Entry]) -> Result<()> {
         for entry in entries {
             let json = serde_json::to_string(&entry)?;
-            let payload = format!(r#"{{"dt":"{}","entry":{}}}"#, Utc::now().to_rfc3339(), json);
+            let (slot, slot_leader) = {
+                let leader_scheduler = self.leader_scheduler.read().unwrap();
+                let slot = leader_scheduler.tick_height_to_slot(entry.tick_height);
+                (slot, leader_scheduler.get_leader_for_slot(slot))
+            };
+            let payload = format!(
+                r#"{{"dt":"{}","t":"entry","s":{},"leader_id":"{:?}","entry":{}}}"#,
+                Utc::now().to_rfc3339_opts(SecondsFormat::Nanos, true),
+                slot,
+                slot_leader,
+                json
+            );
             self.socket.push(payload);
         }
         Ok(())