Add load and flush during startup/shutdown

Author: markstory

src/main.rs

@@ -87,6 +87,10 @@ async fn main() -> Result<(), Error> {
             Err(err) => error!("Failed to run full vacuum on startup: {:?}", err),
         }
     }
+    info!("Loading metadata state from sqlite");
+    store.load_metadata().await?;
+    info!("Loading metadata complete");
+
     // Get startup time after migrations and vacuum
     let startup_time = Utc::now();
 
@@ -236,5 +240,9 @@ async fn main() -> Result<(), Error> {
         .on_completion(log_task_completion("maintenance_task", maintenance_task))
         .await;
 
+    info!("Flushing metadata to sqlite");
+    store.flush_metadata().await?;
+    info!("Shutdown complete");
+
     Ok(())
 }

src/store/inflight_activation.rs

@@ -372,6 +372,28 @@ impl InflightActivationStore {
         )))
     }
 
+    /// Load state into the metadata_store from sqlite.
+    /// Used during application startup to rebuild in-memory state.
+    #[instrument(skip_all)]
+    pub async fn load_metadata(&self) -> Result<(), Error> {
+        self.metadata_store
+            .lock()
+            .await
+            .load_from_sqlite(&self.read_pool)
+            .await
+    }
+
+    /// Flush any pending state in metadata_store into sqlite
+    pub async fn flush_metadata(&self) -> Result<(), Error> {
+        let atomic = self.write_pool.begin().await?;
+        let res = self.metadata_store.lock().await.commit(atomic).await;
+
+        match res {
+            Ok(_) => Ok(()),
+            Err(err) => Err(err.into()),
+        }
+    }
+
     #[instrument(skip_all)]
     pub async fn store(&self, batch: Vec<InflightActivation>) -> Result<QueryResult, Error> {
         if batch.is_empty() {
@@ -383,66 +405,9 @@ impl InflightActivationStore {
             .map(ActivationMetadata::try_from)
             .collect::<Result<Vec<ActivationMetadata>, _>>()?;
 
-        let rows = batch
-            .clone()
-            .into_iter()
-            .map(TableRow::try_from)
-            .collect::<Result<Vec<TableRow>, _>>()?;
-
-        let mut query_builder = QueryBuilder::<Sqlite>::new(
-            "
-            INSERT INTO inflight_taskactivations
-                (
-                    id,
-                    activation,
-                    partition,
-                    offset,
-                    added_at,
-                    received_at,
-                    processing_attempts,
-                    expires_at,
-                    delay_until,
-                    processing_deadline_duration,
-                    processing_deadline,
-                    status,
-                    at_most_once,
-                    namespace,
-                    taskname,
-                    on_attempts_exceeded
-                )
-            ",
-        );
-        let query = query_builder
-            .push_values(rows, |mut b, row| {
-                b.push_bind(row.id);
-                b.push_bind(row.activation);
-                b.push_bind(row.partition);
-                b.push_bind(row.offset);
-                b.push_bind(row.added_at.timestamp());
-                b.push_bind(row.received_at.timestamp());
-                b.push_bind(row.processing_attempts);
-                b.push_bind(row.expires_at.map(|t| Some(t.timestamp())));
-                b.push_bind(row.delay_until.map(|t| Some(t.timestamp())));
-                b.push_bind(row.processing_deadline_duration);
-                if let Some(deadline) = row.processing_deadline {
-                    b.push_bind(deadline.timestamp());
-                } else {
-                    // Add a literal null
-                    b.push("null");
-                }
-                b.push_bind(row.status);
-                b.push_bind(row.at_most_once);
-                b.push_bind(row.namespace);
-                b.push_bind(row.taskname);
-                b.push_bind(row.on_attempts_exceeded as i32);
-            })
-            .push(" ON CONFLICT(id) DO NOTHING")
-            .build();
         let mut atomic = self.write_pool.begin().await?;
-        let meta_result = Ok(query.execute(&mut *atomic).await?.into());
 
-        // insert into the separate stores.
-        // TODO these queries should use one loop.
+        // Insert into the blob store and metadata
         let mut query_builder =
             QueryBuilder::<Sqlite>::new("INSERT INTO activation_blobs (id, activation) ");
         let query = query_builder
@@ -452,7 +417,7 @@ impl InflightActivationStore {
             })
             .push(" ON CONFLICT(id) DO NOTHING")
             .build();
-        query.execute(&mut *atomic).await?;
+        let blob_result = query.execute(&mut *atomic).await?;
 
         {
             // append metadata to memory store and flush to sqlite.
@@ -470,7 +435,7 @@ impl InflightActivationStore {
         metrics::gauge!("store.pages_committed_to_db")
             .set(checkpoint_result.get::<i32, _>("checkpointed"));
 
-        meta_result
+        Ok(blob_result.into())
     }
 
     #[instrument(skip_all)]
@@ -593,10 +558,6 @@ impl InflightActivationStore {
     #[instrument(skip_all)]
     pub async fn delete_activation(&self, id: &str) -> Result<(), Error> {
         self.metadata_store.lock().await.delete(id);
-        sqlx::query("DELETE FROM inflight_taskactivations WHERE id = $1")
-            .bind(id)
-            .execute(&self.write_pool)
-            .await?;
         sqlx::query("DELETE FROM activation_blobs WHERE id = $1")
             .bind(id)
             .execute(&self.write_pool)
@@ -642,9 +603,6 @@ impl InflightActivationStore {
 
     pub async fn clear(&self) -> Result<(), Error> {
         let mut atomic = self.write_pool.begin().await?;
-        sqlx::query("DELETE FROM inflight_taskactivations")
-            .execute(&mut *atomic)
-            .await?;
         sqlx::query("DELETE FROM activation_blobs")
             .execute(&mut *atomic)
             .await?;
@@ -679,24 +637,8 @@ impl InflightActivationStore {
     /// These tasks are set to status=failure and will be handled by handle_failed_tasks accordingly.
     #[instrument(skip_all)]
     pub async fn handle_processing_attempts(&self) -> Result<u64, Error> {
-        // TODO remove this method? It is a no-op with metadata_store
-        // as processing attempts are handled in get_pending_activation
-        let processing_attempts_result = sqlx::query(
-            "UPDATE inflight_taskactivations 
-            SET status = $1 
-            WHERE processing_attempts >= $2 AND status = $3",
-        )
-        .bind(InflightActivationStatus::Failure)
-        .bind(self.config.max_processing_attempts as i32)
-        .bind(InflightActivationStatus::Pending)
-        .execute(&self.write_pool)
-        .await;
-
-        if let Ok(query_res) = processing_attempts_result {
-            return Ok(query_res.rows_affected());
-        }
-
-        Err(anyhow!("Could not update tasks past processing_deadline"))
+        // TODO This is a no-op with metadata_store
+        Ok(0)
     }
 
     /// Perform upkeep work for tasks that are past expires_at deadlines
@@ -707,17 +649,8 @@ impl InflightActivationStore {
     /// The number of impacted records is returned in a Result.
     #[instrument(skip_all)]
     pub async fn handle_expires_at(&self) -> Result<u64, Error> {
-        // TODO: Remove this? This is a no-op with the metadata store.
-        let now = Utc::now();
-        let query = sqlx::query(
-            "DELETE FROM inflight_taskactivations WHERE status = $1 AND expires_at IS NOT NULL AND expires_at < $2",
-        )
-        .bind(InflightActivationStatus::Pending)
-        .bind(now.timestamp())
-        .execute(&self.write_pool)
-        .await?;
-
-        Ok(query.rows_affected())
+        // TODO: This is a no-op with the metadata store.
+        Ok(0)
     }
 
     /// Perform upkeep work for tasks that are past delay_until deadlines

src/store/metadata_store.rs

@@ -6,7 +6,7 @@ use std::{
 use chrono::{DateTime, Duration, Utc};
 use sentry::capture_message;
 use sentry_protos::taskbroker::v1::OnAttemptsExceeded;
-use sqlx::{QueryBuilder, Sqlite, Transaction, sqlite::SqliteQueryResult};
+use sqlx::{QueryBuilder, Sqlite, SqlitePool, Transaction, sqlite::SqliteQueryResult};
 
 use crate::store::inflight_activation::InflightActivationStatus;
 use crate::store::records::{ActivationMetadata, TimestampEntry};
@@ -78,6 +78,20 @@ impl MetadataStore {
         }
     }
 
+    /// Load metadata state from sqlite into memory.
+    pub async fn load_from_sqlite(&mut self, connection: &SqlitePool) -> anyhow::Result<()> {
+        let load_query: Vec<ActivationMetadata> =
+            sqlx::query_as("SELECT * FROM activation_metadata")
+                .fetch_all(connection)
+                .await?
+                .into_iter()
+                .collect();
+        for row in load_query.into_iter() {
+            self.upsert(row)?;
+        }
+        Ok(())
+    }
+
     /// Insert or update an activation metadata record
     pub fn upsert(&mut self, metadata: ActivationMetadata) -> anyhow::Result<()> {
         // File by status
@@ -377,7 +391,6 @@ impl MetadataStore {
                 self.add_to_failed(entry.id.clone());
                 continue;
             }
-            // TODO: Check if the entry is killed
 
             // Filter by namespace if provided.
             if let Some(namespace) = namespace {