Merge pull request #757 from obeli-sk/fix-deployment-id-after-failed-upgrade

fix(db): Update deployment id only after successful auto-upgrade

Author: tomasol

crates/db-postgres/src/postgres_dao.rs

@@ -908,7 +908,7 @@ async fn update_state_component_upgrade_finished_failed(
 async fn update_state_locked_get_intermittent_event_count(
     tx: &Transaction<'_>,
     execution_id: &ExecutionId,
-    deployment_id: DeploymentId,
+    deployment_id: Option<DeploymentId>,
     component_digest: Option<&ComponentDigest>,
     executor_id: ExecutorId,
     run_id: RunId,
@@ -939,7 +939,7 @@ async fn update_state_locked_get_intermittent_event_count(
                 pending_expires_finished = $2,
                 state = $3,
                 updated_at = CURRENT_TIMESTAMP,
-                deployment_id = $4,
+                deployment_id = COALESCE($4, deployment_id),
                 component_id_input_digest = COALESCE($5, component_id_input_digest),
 
                 max_retries = $6,
@@ -959,7 +959,7 @@ async fn update_state_locked_get_intermittent_event_count(
                 &i64::from(appending_version.0),
                 &lock_expires_at,
                 &STATE_LOCKED,
-                &deployment_id.to_string(),
+                &deployment_id.map(|deployment_id| deployment_id.to_string()),
                 &component_digest.map(|component_digest| component_digest.as_slice().to_vec()), // no change if `None` due to `coalesce`
                 &retry_config.max_retries.map(i64::from),
                 &backoff_millis,
@@ -2029,7 +2029,7 @@ fn parse_response_with_cursor(
 }
 
 /// `component_id` is used to construct the `Locked` event.
-/// If `update_component_digest` is set, `t_state` will be set to the component's digest.
+/// If `update_component_digest_and_deployment_id` is set, execution state (`t_state`) will be updated to the component's digest and deployment id.
 /// This should be true in all cases except for `lock_pending_by_ffqns_auto`, where
 /// the digest in `t_state` is only updated after a successful execution upgrade.
 #[instrument(level = Level::TRACE, skip_all, fields(%execution_id, %run_id, %executor_id))]
@@ -2038,7 +2038,7 @@ async fn lock_single_execution(
     tx: &Transaction<'_>,
     created_at: DateTime<Utc>,
     component_id: &ComponentId,
-    update_component_digest: bool, // if set, update `t_state` as well
+    update_component_digest_and_deployment_id: bool,
     deployment_id: DeploymentId,
     execution_id: &ExecutionId,
     run_id: RunId,
@@ -2051,7 +2051,7 @@ async fn lock_single_execution(
 
     // Check State
     let combined_state = get_combined_state(tx, execution_id).await?;
-    let context_component_digest = if update_component_digest {
+    let context_component_digest = if update_component_digest_and_deployment_id {
         component_id.component_digest.clone()
     } else {
         combined_state.execution_with_state.component_digest.clone()
@@ -2106,8 +2106,8 @@ async fn lock_single_execution(
     let intermittent_event_count = update_state_locked_get_intermittent_event_count(
         tx,
         execution_id,
-        deployment_id,
-        update_component_digest.then_some(&component_id.component_digest),
+        update_component_digest_and_deployment_id.then_some(deployment_id),
+        update_component_digest_and_deployment_id.then_some(&component_id.component_digest),
         executor_id,
         run_id,
         lock_expires_at,

crates/db-sqlite/src/sqlite_dao.rs

@@ -1214,7 +1214,7 @@ impl SqlitePool {
     fn update_state_locked_get_intermittent_event_count(
         tx: &Transaction,
         execution_id: &ExecutionId,
-        deployment_id: DeploymentId,
+        deployment_id: Option<DeploymentId>,
         component_digest: Option<&ComponentDigest>,
         executor_id: ExecutorId,
         run_id: RunId,
@@ -1241,7 +1241,7 @@ impl SqlitePool {
                     pending_expires_finished = :pending_expires_finished,
                     state = :state,
                     updated_at = CURRENT_TIMESTAMP,
-                    deployment_id = :deployment_id,
+                    deployment_id = COALESCE(:deployment_id, deployment_id),
                     component_id_input_digest = COALESCE(:component_id_input_digest, component_id_input_digest),
 
                     max_retries = :max_retries,
@@ -1263,7 +1263,7 @@ impl SqlitePool {
             ":appending_version": appending_version.0,
             ":pending_expires_finished": lock_expires_at,
             ":state": STATE_LOCKED,
-            ":deployment_id": deployment_id.to_string(),
+            ":deployment_id": deployment_id, // no change if `None` due to `coalesce`
             ":component_id_input_digest": component_digest.cloned(), // no change if `None` due to `coalesce`
             ":max_retries": retry_config.max_retries,
             ":retry_exp_backoff_millis": backoff_millis,
@@ -1852,7 +1852,7 @@ impl SqlitePool {
     }
 
     /// `component_id` is used to construct the `Locked` event.
-    /// If `update_component_digest` is set, `t_state` will be set to the component's digest.
+    /// If `update_component_digest_and_deployment_id` is set, execution state (`t_state`) will be updated to the component's digest and deployment id.
     /// This should be true in all cases except for `lock_pending_by_ffqns_auto`, where
     /// the digest in `t_state` is only updated after a successful execution upgrade.
     #[instrument(level = Level::TRACE, skip(tx))]
@@ -1861,7 +1861,7 @@ impl SqlitePool {
         tx: &Transaction,
         created_at: DateTime<Utc>,
         component_id: &ComponentId,
-        update_component_digest: bool, // if set, update `t_state` as well
+        update_component_digest_and_deployment_id: bool,
         deployment_id: DeploymentId,
         execution_id: &ExecutionId,
         run_id: RunId,
@@ -1872,7 +1872,7 @@ impl SqlitePool {
     ) -> Result<LockedExecution, DbErrorWrite> {
         trace!("lock_single_execution");
         let combined_state = Self::get_combined_state(tx, execution_id)?;
-        let context_component_digest = if update_component_digest {
+        let context_component_digest = if update_component_digest_and_deployment_id {
             component_id.component_digest.clone()
         } else {
             combined_state.execution_with_state.component_digest.clone()
@@ -1934,8 +1934,8 @@ impl SqlitePool {
         let intermittent_event_count = Self::update_state_locked_get_intermittent_event_count(
             tx,
             execution_id,
-            deployment_id,
-            update_component_digest.then_some(&component_id.component_digest),
+            update_component_digest_and_deployment_id.then_some(deployment_id),
+            update_component_digest_and_deployment_id.then_some(&component_id.component_digest),
             executor_id,
             run_id,
             lock_expires_at,

crates/wasm-workers/src/workflow/workflow_js_worker.rs

@@ -597,7 +597,7 @@ mod tests {
     use assert_matches::assert_matches;
     use chrono::DateTime;
     use concepts::component_id::{COMPONENT_DIGEST_DUMMY, ComponentDigest, Digest};
-    use concepts::prefixed_ulid::{DEPLOYMENT_ID_DUMMY, DelayId, ExecutorId, RunId};
+    use concepts::prefixed_ulid::{DEPLOYMENT_ID_DUMMY, DelayId, DeploymentId, ExecutorId, RunId};
     use concepts::storage::{
         CapturedDbWrite, ComponentUpgradeOutcome, ComponentUpgradeReason, CreateRequest,
         DbConnectionTest, DbPool, DbPoolCloseable, ExecutionRequest, HistoryEvent, JoinSetRequest,
@@ -1059,6 +1059,26 @@ mod tests {
         clock_fn: Box<dyn ClockFn>,
         fn_registry: Arc<dyn FunctionRegistry>,
         workflow_engine: Arc<Engine>,
+    ) -> (WorkflowJsWorker, concepts::ComponentId, RunnableComponent) {
+        compile_js_workflow_worker_with_deployment_id(
+            js_source,
+            user_ffqn,
+            db_pool,
+            clock_fn,
+            fn_registry,
+            workflow_engine,
+            DEPLOYMENT_ID_DUMMY,
+        )
+    }
+
+    fn compile_js_workflow_worker_with_deployment_id(
+        js_source: &str,
+        user_ffqn: &FunctionFqn,
+        db_pool: Arc<dyn DbPool>,
+        clock_fn: Box<dyn ClockFn>,
+        fn_registry: Arc<dyn FunctionRegistry>,
+        workflow_engine: Arc<Engine>,
+        deployment_id: DeploymentId,
     ) -> (WorkflowJsWorker, concepts::ComponentId, RunnableComponent) {
         let wasm_path = workflow_js_runtime_builder::WORKFLOW_JS_RUNTIME;
         let params = default_js_params();
@@ -1108,7 +1128,7 @@ mod tests {
 
         (
             linked.into_worker(
-                DEPLOYMENT_ID_DUMMY,
+                deployment_id,
                 db_pool,
                 deadline_factory,
                 CancelRegistry::new(),
@@ -2575,6 +2595,8 @@ mod tests {
         let (_guard, db_pool, db_close) = database.set_up().await;
         let sim_clock = SimClock::epoch();
         let user_ffqn = FunctionFqn::new_static("test:pkg/ifc", "test-auto-locking-failure");
+        let original_deployment_id = DeploymentId::from_parts(0, 9009);
+        let upgrade_deployment_id = DeploymentId::from_parts(0, 9010);
         let original_js_source = r"
         export default function test_auto_locking_failure(params) {
             obelisk.sleep({ milliseconds: 10 });
@@ -2594,22 +2616,25 @@ mod tests {
         let workflow_engine =
             Engines::get_workflow_engine_test(EngineConfig::on_demand_testing()).unwrap();
         let (original_worker, original_component_id, _original_runnable) =
-            compile_js_workflow_worker(
+            compile_js_workflow_worker_with_deployment_id(
                 original_js_source,
                 &user_ffqn,
                 db_pool.clone(),
                 sim_clock.clone_box(),
                 fn_registry.clone(),
                 workflow_engine.clone(),
+                original_deployment_id,
+            );
+        let (upgrade_worker, upgrade_component_id, _upgrade_runnable) =
+            compile_js_workflow_worker_with_deployment_id(
+                failing_upgrade_js_source,
+                &user_ffqn,
+                db_pool.clone(),
+                sim_clock.clone_box(),
+                fn_registry,
+                workflow_engine,
+                upgrade_deployment_id,
             );
-        let (upgrade_worker, upgrade_component_id, _upgrade_runnable) = compile_js_workflow_worker(
-            failing_upgrade_js_source,
-            &user_ffqn,
-            db_pool.clone(),
-            sim_clock.clone_box(),
-            fn_registry,
-            workflow_engine,
-        );
         assert_ne!(
             original_component_id.component_digest,
             upgrade_component_id.component_digest
@@ -2644,8 +2669,8 @@ mod tests {
                 parent: None,
                 metadata: ExecutionMetadata::empty(),
                 scheduled_at: created_at,
-                component_id: original_component_id,
-                deployment_id: DEPLOYMENT_ID_DUMMY,
+                component_id: original_component_id.clone(),
+                deployment_id: original_deployment_id,
                 scheduled_by: None,
                 paused: false,
             })
@@ -2677,6 +2702,8 @@ mod tests {
         );
 
         let log = db_connection.get(&execution_id).await.unwrap();
+        assert_eq!(original_deployment_id, log.deployment_id);
+        assert_eq!(original_component_id.component_digest, log.component_digest);
         assert_eq!(8, log.events.len());
         assert_matches!(&log.events[0].event, ExecutionRequest::Created { .. });
         assert_matches!(&log.events[1].event, ExecutionRequest::Locked(_));