Introduce event streaming into a log file

Author: Kobzol

Cargo.lock

@@ -46,6 +46,7 @@ tui = { version = "0.16" }
 termion = "1.5"
 indicatif = "0.16.2"
 textwrap = "0.14"
+async-compression = { version = "0.3.8", features = ["tokio", "gzip"] }
 
 # Tako
 tako = { path = "../tako" }

crates/hyperqueue/Cargo.toml

@@ -143,6 +143,10 @@ struct ServerStartOpts {
     /// The maximum number of events tako server will store in memory
     #[clap(long, default_value = "1000000")]
     event_store_size: usize,
+
+    /// Path to a log file where events will be stored.
+    #[clap(long, hide(true))]
+    event_log_file: Option<PathBuf>,
 }
 
 #[derive(Parser)]
@@ -292,7 +296,8 @@ async fn command_server_start(
         autoalloc_interval: opts.autoalloc_interval.map(|x| x.unpack()),
         client_port: opts.client_port,
         worker_port: opts.worker_port,
-        event_store_size: opts.event_store_size,
+        event_buffer_size: opts.event_store_size,
+        event_log_file: opts.event_log_file,
     };
 
     init_hq_server(gsettings, server_cfg).await

crates/hyperqueue/src/bin/hq.rs

@@ -0,0 +1,24 @@
+mod stream;
+mod write;
+
+pub use stream::{start_event_streaming, EventStreamSender};
+pub use write::EventLogWriter;
+
+use bstr::BString;
+use serde::{Deserialize, Serialize};
+
+const HQ_LOG_HEADER: &[u8] = b"hq-event-log";
+const HQ_LOG_VERSION: u32 = 0;
+
+fn canonical_header() -> LogFileHeader {
+    LogFileHeader {
+        header: HQ_LOG_HEADER.into(),
+        version: HQ_LOG_VERSION,
+    }
+}
+
+#[derive(Serialize, Deserialize, Debug, Eq, PartialEq)]
+struct LogFileHeader {
+    header: BString,
+    version: u32,
+}

crates/hyperqueue/src/event/log/mod.rs

@@ -0,0 +1,70 @@
+use crate::event::log::write::EventLogWriter;
+use crate::event::MonitoringEvent;
+use std::future::Future;
+use std::time::Duration;
+use tokio::sync::mpsc;
+
+pub type EventStreamSender = mpsc::UnboundedSender<MonitoringEvent>;
+pub type EventStreamReceiver = mpsc::UnboundedReceiver<MonitoringEvent>;
+
+fn create_event_stream_queue() -> (EventStreamSender, EventStreamReceiver) {
+    mpsc::unbounded_channel()
+}
+
+/// Start event streaming into a log file.
+/// Streaming is running on another thread to reduce overhead and interference.
+///
+/// Returns a future that resolves once the event streaming thread finishes.
+/// The thread will finish if there is some I/O error or if the `receiver` is closed.
+pub fn start_event_streaming(
+    writer: EventLogWriter,
+) -> (EventStreamSender, impl Future<Output = ()>) {
+    let (tx, rx) = create_event_stream_queue();
+
+    let handle = std::thread::spawn(move || {
+        let process = streaming_process(writer, rx);
+
+        let runtime = tokio::runtime::Builder::new_current_thread()
+            .enable_all()
+            .build()
+            .unwrap();
+
+        if let Err(error) = runtime.block_on(process) {
+            log::error!("Event streaming has ended with an error: {error:?}");
+        } else {
+            log::debug!("Event streaming has finished successfully");
+        }
+    });
+    let end_fut = async move {
+        handle.join().expect("Event streaming thread has crashed");
+    };
+    (tx, end_fut)
+}
+
+const FLUSH_PERIOD: Duration = Duration::from_secs(30);
+
+async fn streaming_process(
+    mut writer: EventLogWriter,
+    mut receiver: EventStreamReceiver,
+) -> anyhow::Result<()> {
+    let mut flush_fut = tokio::time::interval(FLUSH_PERIOD);
+
+    loop {
+        tokio::select! {
+            _ = flush_fut.tick() => {
+                writer.flush().await?;
+            }
+            res = receiver.recv() => {
+                match res {
+                    Some(event) => {
+                        log::trace!("Event: {event:?}");
+                        writer.store(event).await?;
+                    }
+                    None => break
+                }
+            }
+        }
+    }
+    writer.finish().await?;
+    Ok(())
+}

crates/hyperqueue/src/event/log/stream.rs

@@ -0,0 +1,65 @@
+use crate::event::log::canonical_header;
+use crate::event::MonitoringEvent;
+use async_compression::tokio::write::GzipEncoder;
+use async_compression::Level;
+use std::path::Path;
+use tokio::fs::File;
+use tokio::io::AsyncWriteExt;
+
+/// Streams monitoring events into a file on disk.
+pub struct EventLogWriter {
+    file: GzipEncoder<File>,
+    buffer: Vec<u8>,
+}
+
+const BUF_MAX_SIZE: usize = 16 * 1024;
+
+impl EventLogWriter {
+    pub async fn create(path: &Path) -> anyhow::Result<Self> {
+        let mut file = File::create(path).await?;
+        let header = rmp_serde::encode::to_vec(&canonical_header())?;
+        file.write_all(&header).await?;
+        file.flush().await?;
+
+        let file = GzipEncoder::with_quality(file, Level::Fastest);
+
+        // Keep buffer capacity larger than max size to avoid reallocation if we overflow
+        // the buffer.
+        let buffer = Vec::with_capacity(BUF_MAX_SIZE * 2);
+        Ok(Self { file, buffer })
+    }
+
+    #[inline]
+    pub async fn store(&mut self, event: MonitoringEvent) -> anyhow::Result<()> {
+        rmp_serde::encode::write(&mut self.buffer, &event)?;
+        if self.is_buffer_full() {
+            self.write_buffer().await?;
+        }
+        Ok(())
+    }
+
+    pub async fn flush(&mut self) -> anyhow::Result<()> {
+        if !self.buffer.is_empty() {
+            self.write_buffer().await?;
+        }
+        self.file.flush().await?;
+        Ok(())
+    }
+
+    pub async fn finish(mut self) -> anyhow::Result<()> {
+        self.flush().await?;
+        self.file.shutdown().await?;
+        Ok(())
+    }
+
+    async fn write_buffer(&mut self) -> tokio::io::Result<()> {
+        self.file.write_all(&self.buffer).await?;
+        self.buffer.clear();
+        Ok(())
+    }
+
+    #[inline]
+    fn is_buffer_full(&self) -> bool {
+        self.buffer.len() >= BUF_MAX_SIZE
+    }
+}

crates/hyperqueue/src/event/log/write.rs

@@ -1,4 +1,5 @@
 pub mod events;
+pub mod log;
 pub mod storage;
 
 use events::MonitoringEventPayload;

crates/hyperqueue/src/event/mod.rs

@@ -1,4 +1,5 @@
 use crate::event::events::MonitoringEventPayload;
+use crate::event::log::EventStreamSender;
 use crate::event::{MonitoringEvent, MonitoringEventId};
 use crate::WorkerId;
 use std::collections::VecDeque;
@@ -11,6 +12,7 @@ pub struct EventStorage {
     event_store_size: usize,
     event_queue: VecDeque<MonitoringEvent>,
     last_event_id: u32,
+    stream_sender: Option<EventStreamSender>,
 }
 
 impl Default for EventStorage {
@@ -19,16 +21,18 @@ impl Default for EventStorage {
             event_store_size: 1_000_000,
             event_queue: Default::default(),
             last_event_id: 0,
+            stream_sender: None,
         }
     }
 }
 
 impl EventStorage {
-    pub fn new(event_store_size: usize) -> Self {
+    pub fn new(event_store_size: usize, stream_sender: Option<EventStreamSender>) -> Self {
         Self {
             event_store_size,
             event_queue: VecDeque::new(),
             last_event_id: 0,
+            stream_sender,
         }
     }
 
@@ -65,13 +69,26 @@ impl EventStorage {
 
     fn insert_event(&mut self, payload: MonitoringEventPayload) {
         self.last_event_id += 1;
-        self.event_queue.push_back(MonitoringEvent {
+
+        let event = MonitoringEvent {
             payload,
             id: self.last_event_id,
             time: SystemTime::now(),
-        });
+        };
+        self.stream_event(&event);
+        self.event_queue.push_back(event);
+
         if self.event_queue.len() > self.event_store_size {
             self.event_queue.pop_front();
         }
     }
+
+    fn stream_event(&mut self, event: &MonitoringEvent) {
+        if let Some(ref streamer) = self.stream_sender {
+            if streamer.send(event.clone()).is_err() {
+                log::error!("Event streaming queue has been closed.");
+                self.stream_sender = None;
+            }
+        }
+    }
 }

crates/hyperqueue/src/event/storage.rs

@@ -20,7 +20,7 @@ pub type Result<T> = std::result::Result<T, Error>;
 // ID types
 use tako::define_id_type;
 
-pub type WorkerId = tako::WorkerId;
+pub use tako::WorkerId;
 pub type TakoTaskId = tako::TaskId;
 pub type Priority = tako::Priority;