Add first version of write-ahead log

Cargo.toml

@@ -39,6 +39,7 @@ prost-build = "0.14.1"
 rand = "0.9.2"
 rustyline = "17.0.2"
 serde = { version = "1.0.228", features = ["derive"] }
+serde_json = "1.0.149"
 snmalloc-rs = "0.3.8"
 sqlparser = "0.59.0"
 sysinfo = "0.37.2"

crates/modelardb_server/src/configuration.rs

@@ -382,6 +382,7 @@ mod tests {
     use std::sync::Arc;
 
     use modelardb_storage::data_folder::DataFolder;
+    use modelardb_storage::write_ahead_log::WriteAheadLog;
     use modelardb_types::types::{Node, ServerMode};
     use tempfile::TempDir;
     use tokio::sync::RwLock;
@@ -689,6 +690,10 @@ mod tests {
             .await
             .unwrap();
 
+        let write_ahead_log = Arc::new(RwLock::new(
+            WriteAheadLog::try_new(&local_data_folder).await.unwrap(),
+        ));
+
         let configuration_manager = Arc::new(RwLock::new(
             ConfigurationManager::try_new(
                 local_data_folder,
@@ -699,7 +704,7 @@ mod tests {
         ));
 
         let storage_engine = Arc::new(RwLock::new(
-            StorageEngine::try_new(data_folders, &configuration_manager)
+            StorageEngine::try_new(data_folders, write_ahead_log, &configuration_manager)
                 .await
                 .unwrap(),
         ));

crates/modelardb_server/src/context.rs

@@ -20,9 +20,10 @@ use std::sync::Arc;
 
 use datafusion::arrow::datatypes::Schema;
 use datafusion::catalog::{SchemaProvider, TableProvider};
+use modelardb_storage::write_ahead_log::WriteAheadLog;
 use modelardb_types::types::TimeSeriesTableMetadata;
 use tokio::sync::RwLock;
-use tracing::info;
+use tracing::{info, warn};
 
 use crate::configuration::ConfigurationManager;
 use crate::error::{ModelarDbServerError, Result};
@@ -38,6 +39,8 @@ pub struct Context {
     pub configuration_manager: Arc<RwLock<ConfigurationManager>>,
     /// Manages all uncompressed and compressed data in the system.
     pub storage_engine: Arc<RwLock<StorageEngine>>,
+    /// Write-ahead log for persisting data and operations.
+    write_ahead_log: Arc<RwLock<WriteAheadLog>>,
 }
 
 impl Context {
@@ -50,14 +53,24 @@ impl Context {
                 .await?,
         ));
 
+        let write_ahead_log = Arc::new(RwLock::new(
+            WriteAheadLog::try_new(&data_folders.local_data_folder).await?,
+        ));
+
         let storage_engine = Arc::new(RwLock::new(
-            StorageEngine::try_new(data_folders.clone(), &configuration_manager).await?,
+            StorageEngine::try_new(
+                data_folders.clone(),
+                write_ahead_log.clone(),
+                &configuration_manager,
+            )
+            .await?,
         ));
 
         Ok(Context {
             data_folders,
             configuration_manager,
             storage_engine,
+            write_ahead_log,
         })
     }
 
@@ -103,6 +116,12 @@ impl Context {
         self.register_and_save_time_series_table(time_series_table_metadata)
             .await?;
 
+        // Create a file in the write-ahead log to log uncompressed data for the table.
+        let mut write_ahead_log = self.write_ahead_log.write().await;
+        write_ahead_log
+            .create_table_log(time_series_table_metadata)
+            .await?;
+
         Ok(())
     }
 
@@ -237,6 +256,40 @@ impl Context {
         Ok(())
     }
 
+    /// For each time series table in the local data folder, use the write-ahead log to replay any
+    /// data that was written to the storage engine but not compressed and saved to disk. Note that
+    /// this method should only be called before the storage engine starts ingesting data to avoid
+    /// replaying data that is currently in memory.
+    pub(super) async fn replay_write_ahead_log(&self) -> Result<()> {
+        let local_data_folder = &self.data_folders.local_data_folder;
+
+        let write_ahead_log = self.write_ahead_log.write().await;
+        let mut storage_engine = self.storage_engine.write().await;
+
+        for metadata in local_data_folder.time_series_table_metadata().await? {
+            let unpersisted_batches =
+                write_ahead_log.unpersisted_batches_in_table_log(&metadata.name)?;
+
+            if !unpersisted_batches.is_empty() {
+                warn!(
+                    table = %metadata.name,
+                    batch_count = unpersisted_batches.len(),
+                    "Replaying unpersisted batches for time series table."
+                );
+            }
+
+            for (batch_id, batch) in unpersisted_batches {
+                storage_engine.insert_data_points_with_batch_id(
+                    metadata.clone(),
+                    batch,
+                    batch_id,
+                )?;
+            }
+        }
+
+        Ok(())
+    }
+
     /// Drop the table with `table_name` if it exists. The table is deregistered from the Apache
     /// Arrow Datafusion session context and deleted from the storage engine and Delta Lake. If the
     /// table does not exist or if it could not be dropped, [`ModelarDbServerError`] is returned.
@@ -254,11 +307,16 @@ impl Context {
 
         self.drop_table_from_storage_engine(table_name).await?;
 
+        let local_data_folder = &self.data_folders.local_data_folder;
+
+        // If the table is a time series table, delete the table log file from the write-ahead log.
+        if local_data_folder.is_time_series_table(table_name).await? {
+            let mut write_ahead_log = self.write_ahead_log.write().await;
+            write_ahead_log.remove_table_log(table_name)?;
+        }
+
         // Drop the table from the Delta Lake.
-        self.data_folders
-            .local_data_folder
-            .drop_table(table_name)
-            .await?;
+        local_data_folder.drop_table(table_name).await?;
 
         Ok(())
     }

crates/modelardb_server/src/main.rs

@@ -86,13 +86,8 @@ async fn main() -> Result<()> {
     // Setup CTRL+C handler.
     setup_ctrl_c_handler(&context);
 
-    // Initialize storage engine with spilled buffers.
-    context
-        .storage_engine
-        .read()
-        .await
-        .initialize(&context)
-        .await?;
+    // Replay any data that was written to the storage engine but not compressed and saved to disk.
+    context.replay_write_ahead_log().await?;
 
     // Start the Apache Arrow Flight interface.
     remote::start_apache_arrow_flight_server(context, *PORT).await?;

crates/modelardb_server/src/storage/compressed_data_buffer.rs

@@ -15,6 +15,7 @@
 
 //! Buffer for compressed segments from the same time series table.
 
+use std::collections::HashSet;
 use std::sync::Arc;
 
 use datafusion::arrow::record_batch::RecordBatch;
@@ -30,16 +31,21 @@ pub(super) struct CompressedSegmentBatch {
     pub(super) time_series_table_metadata: Arc<TimeSeriesTableMetadata>,
     /// Compressed segments representing the data points to insert.
     pub(super) compressed_segments: Vec<RecordBatch>,
+    /// The ids of the uncompressed batches that correspond to the compressed segments. The ids are
+    /// assigned by the WAL and are used to delete uncompressed data when compressed data is saved.
+    pub(super) batch_ids: HashSet<u64>,
 }
 
 impl CompressedSegmentBatch {
     pub(super) fn new(
         time_series_table_metadata: Arc<TimeSeriesTableMetadata>,
         compressed_segments: Vec<RecordBatch>,
+        batch_ids: HashSet<u64>,
     ) -> Self {
         Self {
             time_series_table_metadata,
             compressed_segments,
+            batch_ids,
         }
     }
 
@@ -59,6 +65,9 @@ pub(super) struct CompressedDataBuffer {
     compressed_segments: Vec<RecordBatch>,
     /// Continuously updated total sum of the size of the compressed segments.
     pub(super) size_in_bytes: u64,
+    /// The ids of the uncompressed batches that correspond to the compressed segments. The ids are
+    /// assigned by the WAL and are used to delete uncompressed data when compressed data is saved.
+    batch_ids: HashSet<u64>,
 }
 
 impl CompressedDataBuffer {
@@ -67,16 +76,19 @@ impl CompressedDataBuffer {
             time_series_table_metadata,
             compressed_segments: vec![],
             size_in_bytes: 0,
+            batch_ids: HashSet::new(),
         }
     }
 
-    /// Append `compressed_segments` to the [`CompressedDataBuffer`] and return the size of
-    /// `compressed_segments` in bytes if their schema matches the time series table, otherwise
-    /// [`ModelarDbServerError`] is returned.
-    pub(super) fn append_compressed_segments(
+    /// Append the compressed segments in `compressed_segment_batch` to the [`CompressedDataBuffer`]
+    /// and return the size of the compressed segments in bytes if their schema matches the time
+    /// series table, otherwise [`ModelarDbServerError`] is returned.
+    pub(super) fn append_compressed_segment_batch(
         &mut self,
-        mut compressed_segments: Vec<RecordBatch>,
+        compressed_segment_batch: CompressedSegmentBatch,
     ) -> Result<u64> {
+        let mut compressed_segments = compressed_segment_batch.compressed_segments;
+
         if compressed_segments.iter().any(|compressed_segments| {
             compressed_segments.schema() != self.time_series_table_metadata.compressed_schema
         }) {
@@ -94,6 +106,8 @@ impl CompressedDataBuffer {
             self.size_in_bytes += compressed_segments_size;
         }
 
+        self.batch_ids.extend(compressed_segment_batch.batch_ids);
+
         Ok(compressed_segments_size)
     }
 
@@ -102,6 +116,11 @@ impl CompressedDataBuffer {
         self.compressed_segments
     }
 
+    /// Return the ids given to the uncompressed batches by the WAL.
+    pub(super) fn batch_ids(&self) -> HashSet<u64> {
+        self.batch_ids.clone()
+    }
+
     /// Return the size in bytes of `compressed_segments`.
     fn size_of_compressed_segments(compressed_segments: &RecordBatch) -> u64 {
         let mut total_size: u64 = 0;
@@ -128,15 +147,12 @@ mod tests {
     use modelardb_test::table;
 
     #[test]
-    fn test_can_append_valid_compressed_segments() {
+    fn test_can_append_valid_compressed_segment_batch() {
         let mut compressed_data_buffer =
             CompressedDataBuffer::new(table::time_series_table_metadata_arc());
 
         compressed_data_buffer
-            .append_compressed_segments(vec![
-                table::compressed_segments_record_batch(),
-                table::compressed_segments_record_batch(),
-            ])
+            .append_compressed_segment_batch(compressed_segment_batch())
             .unwrap();
 
         assert_eq!(compressed_data_buffer.compressed_segments.len(), 2);
@@ -150,10 +166,7 @@ mod tests {
             CompressedDataBuffer::new(table::time_series_table_metadata_arc());
 
         compressed_data_buffer
-            .append_compressed_segments(vec![
-                table::compressed_segments_record_batch(),
-                table::compressed_segments_record_batch(),
-            ])
+            .append_compressed_segment_batch(compressed_segment_batch())
             .unwrap();
 
         assert!(compressed_data_buffer.size_in_bytes > 0);
@@ -164,12 +177,8 @@ mod tests {
         let mut compressed_data_buffer =
             CompressedDataBuffer::new(table::time_series_table_metadata_arc());
 
-        let compressed_segments = vec![
-            table::compressed_segments_record_batch(),
-            table::compressed_segments_record_batch(),
-        ];
         compressed_data_buffer
-            .append_compressed_segments(compressed_segments)
+            .append_compressed_segment_batch(compressed_segment_batch())
             .unwrap();
 
         let record_batches = compressed_data_buffer.record_batches();
@@ -179,6 +188,17 @@ mod tests {
         assert_eq!(record_batch.num_rows(), 6);
     }
 
+    fn compressed_segment_batch() -> CompressedSegmentBatch {
+        CompressedSegmentBatch::new(
+            table::time_series_table_metadata_arc(),
+            vec![
+                table::compressed_segments_record_batch(),
+                table::compressed_segments_record_batch(),
+            ],
+            HashSet::from([0, 1, 2]),
+        )
+    }
+
     #[test]
     fn test_get_size_of_compressed_data_buffer() {
         let compressed_data_buffer = table::compressed_segments_record_batch();