fix: deletion of invocation impacting whole invocation secondary index

Author: seriousben

server/data_model/src/lib.rs

@@ -830,18 +830,6 @@ impl GraphInvocationCtx {
         key
     }
 
-    pub fn secondary_index_key_prefix_from_compute_graph(
-        namespace: &str,
-        compute_graph_name: &str,
-    ) -> Vec<u8> {
-        let mut key = Vec::new();
-        key.extend_from_slice(namespace.as_bytes());
-        key.push(b'|');
-        key.extend_from_slice(compute_graph_name.as_bytes());
-        key.push(b'|');
-        key
-    }
-
     pub fn get_invocation_id_from_secondary_index_key(key: &[u8]) -> Option<String> {
         key.split(|&b| b == b'|')
             .nth(3)

server/state_store/src/migrations.rs

@@ -18,7 +18,7 @@ use crate::{
     state_machine::IndexifyObjectsColumns,
 };
 
-const SERVER_DB_VERSION: u64 = 7;
+const SERVER_DB_VERSION: u64 = 8;
 
 // Note: should never be used with data model types to guarantee it works with
 // different versions.
@@ -116,6 +116,12 @@ pub fn migrate(path: &Path) -> Result<StateMachineMetadata> {
                 .context("migrating from v6 to v7")?;
         }
 
+        if sm_meta.db_version == 7 {
+            sm_meta.db_version += 1;
+            migrate_v7_to_v8_recompute_invocation_ctx_secondary_index(&db, &txn)
+                .context("migrating from v7 to v8")?;
+        }
+
         // add new migrations before this line and increment SERVER_DB_VERSION
     }
 
@@ -536,6 +542,52 @@ pub fn migrate_v6_to_v7_reallocate_allocated_tasks(
     Ok(())
 }
 
+#[tracing::instrument(skip(db, txn))]
+pub fn migrate_v7_to_v8_recompute_invocation_ctx_secondary_index(
+    db: &TransactionDB,
+    txn: &Transaction<TransactionDB>,
+) -> Result<()> {
+    let mut num_total_invocation_ctx: usize = 0;
+    let mut num_migrated_invocation_ctx: usize = 0;
+
+    panic!("TODO");
+
+    {
+        let mut read_options = ReadOptions::default();
+        read_options.set_readahead_size(4_194_304);
+
+        let iter = db.iterator_cf_opt(
+            &IndexifyObjectsColumns::GraphInvocationCtx.cf_db(&db),
+            read_options,
+            IteratorMode::Start,
+        );
+
+        for kv in iter {
+            num_total_invocation_ctx += 1;
+            let (_key, value) = kv?;
+
+            let graph_invocation_ctx = JsonEncoder::decode::<GraphInvocationCtx>(&value)?;
+
+            let secondary_index_key =
+                GraphInvocationCtx::secondary_index_key(&graph_invocation_ctx);
+
+            txn.put_cf(
+                &IndexifyObjectsColumns::GraphInvocationCtxSecondaryIndex.cf_db(&db),
+                &secondary_index_key,
+                &[],
+            )?;
+
+            num_migrated_invocation_ctx += 1;
+        }
+    }
+
+    info!(
+        "Migrated {}/{} invocation context secondary indexes from v3 to v4",
+        num_migrated_invocation_ctx, num_total_invocation_ctx
+    );
+    Ok(())
+}
+
 pub fn write_sm_meta(
     db: &TransactionDB,
     txn: &Transaction<TransactionDB>,

server/state_store/src/state_machine.rs

@@ -161,14 +161,39 @@ pub(crate) fn delete_invocation(
         "Deleting invocation",
     );
 
-    // Delete the invocation payload
-    let invocation_key =
+    // Check if the invocation was deleted before the task completes
+    let invocation_ctx_key =
         GraphInvocationCtx::key_from(&req.namespace, &req.compute_graph, &req.invocation_id);
-    delete_cf_prefix(
-        txn,
-        &IndexifyObjectsColumns::GraphInvocations.cf_db(&db),
-        invocation_key.as_bytes(),
-    )?;
+    let invocation_ctx = txn
+        .get_cf(
+            &IndexifyObjectsColumns::GraphInvocationCtx.cf_db(&db),
+            &invocation_ctx_key,
+        )
+        .map_err(|e| anyhow!("failed to get invocation: {}", e))?;
+    let invocation_ctx = match invocation_ctx {
+        Some(v) => JsonEncoder::decode::<GraphInvocationCtx>(&v)?,
+        None => {
+            info!(
+                namespace = &req.namespace,
+                compute_graph = &req.compute_graph,
+                invocation_id = &req.invocation_id,
+                "Invocation to delete not found: {}",
+                &req.invocation_id
+            );
+            return Ok(());
+        }
+    };
+
+    // Delete the invocation payload
+    {
+        let invocation_key =
+            InvocationPayload::key_from(&req.namespace, &req.compute_graph, &req.invocation_id);
+
+        txn.delete_cf(
+            &IndexifyObjectsColumns::GraphInvocations.cf_db(&db),
+            &invocation_key,
+        )?;
+    }
 
     let mut tasks_deleted = Vec::new();
     let task_prefix =
@@ -247,26 +272,16 @@ pub(crate) fn delete_invocation(
     }
 
     // Delete Graph Invocation Context
-
     delete_cf_prefix(
         txn,
         IndexifyObjectsColumns::GraphInvocationCtx.cf_db(&db),
-        invocation_key.as_bytes(),
+        invocation_ctx_key.as_bytes(),
     )?;
 
     // Delete Graph Invocation Context Secondary Index
-    // Note We don't delete the secondary index here because it's too much work to
-    // get the invocation id from the secondary index key. We purge all the
-    // secondary index keys for graphs if they are ever deleted.
-    //
-    // TODO: Only delete the secondary index keys for this invocation
-    delete_cf_prefix(
-        txn,
+    txn.delete_cf(
         IndexifyObjectsColumns::GraphInvocationCtxSecondaryIndex.cf_db(&db),
-        &GraphInvocationCtx::secondary_index_key_prefix_from_compute_graph(
-            &req.namespace,
-            &req.compute_graph,
-        ),
+        &invocation_ctx.secondary_index_key(),
     )?;
 
     let node_output_prefix =
@@ -580,13 +595,6 @@ pub fn delete_compute_graph(
         delete_invocation(db.clone(), txn, &req)?;
     }
 
-    // Delete Graph Invocation Context Secondary Index
-    delete_cf_prefix(
-        txn,
-        IndexifyObjectsColumns::GraphInvocationCtxSecondaryIndex.cf_db(&db),
-        &GraphInvocationCtx::secondary_index_key_prefix_from_compute_graph(namespace, name),
-    )?;
-
     for iter in make_prefix_iterator(
         txn,
         &IndexifyObjectsColumns::ComputeGraphVersions.cf_db(&db),