Add support for starting/restarting without state in leader mode

Author: mwylde

crates/arroyo-api/queries/api_queries.sql

@@ -183,6 +183,7 @@ UPDATE job_configs
 SET
    updated_at = :updated_at,
    updated_by = :updated_by,
+   stop = 'none',
    restart_nonce = restart_nonce + 1,
    restart_mode = :mode,
    ignore_state_before_epoch = :ignore_state_before_epoch

crates/arroyo-api/src/pipelines.rs

@@ -52,7 +52,7 @@ use crate::rest_utils::{
 use crate::types::public::{PipelineType, RestartMode, StopMode};
 use crate::udfs::build_udf;
 use crate::{connection_tables, to_micros};
-use arroyo_rpc::config::config;
+use arroyo_rpc::config::{JobControllerMode, config};
 use arroyo_rpc::errors::ErrorDomain;
 use arroyo_types::to_millis;
 use cornucopia_async::{Database, DatabaseSource};
@@ -868,29 +868,44 @@ pub async fn restart_pipeline(
     WithRejection(Json(req), _): WithRejection<Json<PipelineRestart>, ApiError>,
 ) -> Result<Json<Pipeline>, ErrorResp> {
     let auth_data = authenticate(&state.database, bearer_auth).await?;
+
     let db = state.database.client().await?;
 
-    let job_id = api_queries::fetch_get_pipeline_jobs(&db, &auth_data.organization_id, &id)
+    let job = api_queries::fetch_get_pipeline_jobs(&db, &auth_data.organization_id, &id)
         .await?
         .into_iter()
         .next()
-        .ok_or_else(|| bad_request("No jobs for pipeline"))?
-        .id;
+        .ok_or_else(|| bad_request("No jobs for pipeline"))?;
 
     let mode = if req.force == Some(true) {
         RestartMode::force
     } else {
         RestartMode::safe
     };
 
-    // If user wants to ignore state, query max checkpoint epoch and compute threshold
     let ignore_before_epoch = if req.ignore_state.unwrap_or(false) {
-        api_queries::fetch_max_checkpoint_epoch(&db, &job_id, &auth_data.organization_id)
-            .await?
-            .into_iter()
-            .next()
-            .and_then(|r| r.max_epoch)
-            .map(|max_epoch| max_epoch + 1)
+        match config().job_controller {
+            JobControllerMode::Controller => {
+                // Controller mode uses this as an epoch threshold.
+                api_queries::fetch_max_checkpoint_epoch(&db, &job.id, &auth_data.organization_id)
+                    .await?
+                    .into_iter()
+                    .next()
+                    .and_then(|r| r.max_epoch)
+                    .map(|max_epoch| max_epoch + 1)
+            }
+            JobControllerMode::Worker => {
+                // Leader mode uses this as the generation that should start without state.
+                Some(
+                    job.run_id
+                        .unwrap_or(0)
+                        .max(0)
+                        .checked_add(1)
+                        .and_then(|generation| generation.try_into().ok())
+                        .ok_or_else(|| bad_request("Job generation is too large to restart"))?,
+                )
+            }
+        }
     } else {
         None
     };
@@ -901,7 +916,7 @@ pub async fn restart_pipeline(
         &auth_data.user_id,
         &mode,
         &ignore_before_epoch,
-        &job_id,
+        &job.id,
         &auth_data.organization_id,
     )
     .await?;

crates/arroyo-controller/src/lib.rs

@@ -72,6 +72,8 @@ pub struct JobConfig {
     parallelism_overrides: HashMap<u32, usize>,
     restart_nonce: i32,
     restart_mode: RestartMode,
+    /// Minimum checkpoint epoch in controller mode; generation to start without state in leader
+    /// mode.
     ignore_state_before_epoch: Option<i32>,
     /// Per-job environment variables forwarded to workers at scheduling time.
     env_vars: serde_json::Value,

crates/arroyo-controller/src/states/scheduling.rs

@@ -444,6 +444,20 @@ async fn get_and_register_checkpoint_info_leader<'a>(
         storage_url: ctx.pipeline_info.state_url.clone(),
     };
     let storage_provider = get_storage_provider(&storage_role).await?;
+    let ignore_state_before_generation = ctx
+        .config
+        .ignore_state_before_epoch
+        .and_then(|generation| generation.try_into().ok())
+        .map(Generation);
+
+    if let Some(generation) = ignore_state_before_generation {
+        info!(
+            message = "ignoring state from before leader generation",
+            job_id = *ctx.config.id,
+            ignore_state_before_generation = generation.0,
+            current_generation = ctx.status.generation,
+        );
+    }
 
     let new_gen = initialize_generation(
         storage_provider.as_ref(),
@@ -452,6 +466,7 @@ async fn get_and_register_checkpoint_info_leader<'a>(
             job_id: JobId(ctx.config.id.clone()),
             generation: Generation(ctx.status.generation),
             updated_at: SystemTime::now(),
+            ignore_state_before_generation,
         },
         true,
     )
@@ -717,21 +732,25 @@ impl State for Scheduling {
         let worker_connects = Arc::try_unwrap(worker_connects).unwrap().into_inner();
         let program = api::ArrowProgram::from(ctx.program.clone());
 
-        // Use ignore_state_before_epoch as default so new checkpoints exceed the threshold
-        let default_epoch = ctx
-            .config
-            .ignore_state_before_epoch
-            .filter(|&t| t > 0)
-            .map(|t| {
-                let epoch = (t - 1) as u64;
-                info!(
-                    message = "starting from ignore_state_before_epoch threshold",
-                    job_id = *ctx.config.id,
-                    default_epoch = epoch,
-                );
-                epoch
-            })
-            .unwrap_or(0);
+        // In controller mode this is an epoch threshold. Leader mode uses the same field as a
+        // generation number, so it must not affect the checkpoint epoch.
+        let default_epoch = if leader_mode {
+            0
+        } else {
+            ctx.config
+                .ignore_state_before_epoch
+                .filter(|&t| t > 0)
+                .map(|t| {
+                    let epoch = (t - 1) as u64;
+                    info!(
+                        message = "starting from ignore_state_before_epoch threshold",
+                        job_id = *ctx.config.id,
+                        default_epoch = epoch,
+                    );
+                    epoch
+                })
+                .unwrap_or(0)
+        };
 
         let start_epoch = checkpoint_info
             .as_ref()

crates/arroyo-state-protocol/src/lib.rs

@@ -461,6 +461,7 @@ mod tests {
                 job_id: JobId::new("J"),
                 generation: Generation(1),
                 updated_at: from_micros(123),
+                ignore_state_before_generation: None,
             },
             false,
         )
@@ -516,6 +517,7 @@ mod tests {
                 job_id: JobId::new("J"),
                 generation: Generation(2),
                 updated_at: from_micros(456),
+                ignore_state_before_generation: None,
             },
             false,
         )
@@ -547,6 +549,75 @@ mod tests {
         assert_eq!(written_manifest, expected_manifest);
     }
 
+    #[tokio::test]
+    async fn initialize_generation_keeps_ignoring_old_state_on_retry() {
+        let store = MemoryProtocolStore::default();
+        let paths = ProtocolPaths::new(PipelineId::new("P"), JobId::new("J"));
+        write_current_generation(&store, &paths, Generation(3)).await;
+
+        let checkpoint_ref = paths.checkpoint_manifest(Generation(1), Epoch(1));
+        let checkpoint = checkpoint_for_generation(Generation(1), 1, None, false);
+        write_canonical_checkpoint(&store, &paths, &checkpoint_ref, &checkpoint).await;
+        let previous_manifest =
+            generation_manifest_for_generation(Generation(1), None, Some(checkpoint_ref));
+        put_json(
+            &store,
+            &paths.generation_manifest(Generation(1)),
+            &previous_manifest,
+        )
+        .await
+        .unwrap();
+        put_json(
+            &store,
+            &paths.generation_manifest(Generation(2)),
+            &GenerationManifest::new(
+                PipelineId::new("P"),
+                JobId::new("J"),
+                Generation(2),
+                None,
+                123,
+            ),
+        )
+        .await
+        .unwrap();
+
+        let initialization = initialize_generation(
+            &store,
+            InitializeGenerationRequest {
+                pipeline_id: PipelineId::new("P"),
+                job_id: JobId::new("J"),
+                generation: Generation(3),
+                updated_at: from_micros(456),
+                ignore_state_before_generation: Some(Generation(2)),
+            },
+            false,
+        )
+        .await
+        .unwrap();
+
+        let expected_manifest = GenerationManifest::new(
+            PipelineId::new("P"),
+            JobId::new("J"),
+            Generation(3),
+            None,
+            456,
+        );
+        assert_eq!(
+            initialization,
+            GenerationInitialization::Initialized {
+                generation_manifest: expected_manifest.clone(),
+                recovery: GenerationRecovery::NoCheckpoint,
+            }
+        );
+
+        let written_manifest: GenerationManifest =
+            read_json(&store, &paths.generation_manifest(Generation(3)))
+                .await
+                .unwrap()
+                .expect("new generation manifest should be written");
+        assert_eq!(written_manifest, expected_manifest);
+    }
+
     #[tokio::test]
     async fn initialize_generation_restores_previous_checkpoint_requiring_commit_replay() {
         let store = MemoryProtocolStore::default();
@@ -573,6 +644,7 @@ mod tests {
                 job_id: JobId::new("J"),
                 generation: Generation(2),
                 updated_at: from_micros(456),
+                ignore_state_before_generation: None,
             },
             false,
         )
@@ -621,6 +693,7 @@ mod tests {
                 job_id: JobId::new("J"),
                 generation: Generation(3),
                 updated_at: from_micros(789),
+                ignore_state_before_generation: None,
             },
             false,
         )
@@ -668,6 +741,7 @@ mod tests {
                 job_id: JobId::new("J"),
                 generation: Generation(2),
                 updated_at: from_micros(456),
+                ignore_state_before_generation: None,
             },
             false,
         )
@@ -724,6 +798,7 @@ mod tests {
                 job_id: JobId::new("J"),
                 generation: Generation(3),
                 updated_at: from_micros(456),
+                ignore_state_before_generation: None,
             },
             false,
         )
@@ -789,6 +864,7 @@ mod tests {
                 job_id: JobId::new("J"),
                 generation: Generation(3),
                 updated_at: from_micros(456),
+                ignore_state_before_generation: None,
             },
             false,
         )
@@ -826,6 +902,7 @@ mod tests {
                 job_id: JobId::new("J"),
                 generation: Generation(2),
                 updated_at: from_micros(456),
+                ignore_state_before_generation: None,
             },
             false,
         )

crates/arroyo-state-protocol/src/workflow.rs

@@ -71,6 +71,8 @@ pub struct InitializeGenerationRequest {
     pub job_id: JobId,
     pub generation: Generation,
     pub updated_at: SystemTime,
+    /// Do not restore checkpoints written before this generation.
+    pub ignore_state_before_generation: Option<Generation>,
 }
 
 /// Checkpoint, if any, that a newly initialized generation should restore from.
@@ -259,7 +261,13 @@ where
         });
     }
 
-    let recovery = find_recovery_checkpoint(store, &paths, request.generation).await?;
+    let recovery = find_recovery_checkpoint(
+        store,
+        &paths,
+        request.generation,
+        request.ignore_state_before_generation,
+    )
+    .await?;
     let base_checkpoint_ref = match &recovery {
         RecoverySearch::Found(recovery) => match recovery {
             GenerationRecovery::NoCheckpoint => None,
@@ -307,15 +315,21 @@ async fn find_recovery_checkpoint<S>(
     store: &S,
     paths: &ProtocolPaths,
     generation: Generation,
+    ignore_state_before_generation: Option<Generation>,
 ) -> Result<RecoverySearch, StoreError>
 where
     S: ProtocolStore + ?Sized,
 {
     let Some(previous_generation) = generation.0.checked_sub(1) else {
         return Ok(RecoverySearch::Found(GenerationRecovery::NoCheckpoint));
     };
+    let first_generation = ignore_state_before_generation.unwrap_or(Generation(0)).0;
+
+    if first_generation > previous_generation {
+        return Ok(RecoverySearch::Found(GenerationRecovery::NoCheckpoint));
+    }
 
-    for previous_generation in (0..=previous_generation).rev() {
+    for previous_generation in (first_generation..=previous_generation).rev() {
         let manifest_ref = paths.generation_manifest(Generation(previous_generation));
         let Some(manifest): Option<GenerationManifest> = read_json(store, &manifest_ref).await?
         else {

crates/arroyo-worker/src/job_controller/controller.rs

@@ -160,6 +160,11 @@ impl WorkerJobController {
                 job_id: worker_context.job_id.clone(),
                 generation: Generation(worker_context.generation),
                 updated_at: SystemTime::now(),
+                // When the controller found no parent, do not let the leader independently fall
+                // back to state from an earlier generation.
+                ignore_state_before_generation: parent_ref
+                    .is_none()
+                    .then_some(Generation(worker_context.generation)),
             },
             false,
         )