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,19 @@ 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 = ctx
+        .config
+        .ignore_state_before_epoch
+        .and_then(|generation| generation.try_into().ok())
+        == Some(ctx.status.generation);
+
+    if ignore_state {
+        info!(
+            message = "starting leader generation without state",
+            job_id = *ctx.config.id,
+            generation = ctx.status.generation,
+        );
+    }
 
     let new_gen = initialize_generation(
         storage_provider.as_ref(),
@@ -452,6 +465,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,
         },
         true,
     )
@@ -717,21 +731,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: false,
             },
             false,
         )
@@ -516,6 +517,7 @@ mod tests {
                 job_id: JobId::new("J"),
                 generation: Generation(2),
                 updated_at: from_micros(456),
+                ignore_state: false,
             },
             false,
         )
@@ -547,6 +549,62 @@ mod tests {
         assert_eq!(written_manifest, expected_manifest);
     }
 
+    #[tokio::test]
+    async fn initialize_generation_can_ignore_previous_checkpoint() {
+        let store = MemoryProtocolStore::default();
+        let paths = ProtocolPaths::new(PipelineId::new("P"), JobId::new("J"));
+        write_current_generation(&store, &paths, Generation(2)).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();
+
+        let initialization = initialize_generation(
+            &store,
+            InitializeGenerationRequest {
+                pipeline_id: PipelineId::new("P"),
+                job_id: JobId::new("J"),
+                generation: Generation(2),
+                updated_at: from_micros(456),
+                ignore_state: true,
+            },
+            false,
+        )
+        .await
+        .unwrap();
+
+        let expected_manifest = GenerationManifest::new(
+            PipelineId::new("P"),
+            JobId::new("J"),
+            Generation(2),
+            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(2)))
+                .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 +631,7 @@ mod tests {
                 job_id: JobId::new("J"),
                 generation: Generation(2),
                 updated_at: from_micros(456),
+                ignore_state: false,
             },
             false,
         )
@@ -621,6 +680,7 @@ mod tests {
                 job_id: JobId::new("J"),
                 generation: Generation(3),
                 updated_at: from_micros(789),
+                ignore_state: false,
             },
             false,
         )
@@ -668,6 +728,7 @@ mod tests {
                 job_id: JobId::new("J"),
                 generation: Generation(2),
                 updated_at: from_micros(456),
+                ignore_state: false,
             },
             false,
         )
@@ -724,6 +785,7 @@ mod tests {
                 job_id: JobId::new("J"),
                 generation: Generation(3),
                 updated_at: from_micros(456),
+                ignore_state: false,
             },
             false,
         )
@@ -789,6 +851,7 @@ mod tests {
                 job_id: JobId::new("J"),
                 generation: Generation(3),
                 updated_at: from_micros(456),
+                ignore_state: false,
             },
             false,
         )
@@ -826,6 +889,7 @@ mod tests {
                 job_id: JobId::new("J"),
                 generation: Generation(2),
                 updated_at: from_micros(456),
+                ignore_state: false,
             },
             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,
+    /// Start this generation without restoring a checkpoint from an earlier generation.
+    pub ignore_state: bool,
 }
 
 /// Checkpoint, if any, that a newly initialized generation should restore from.
@@ -259,7 +261,11 @@ where
         });
     }
 
-    let recovery = find_recovery_checkpoint(store, &paths, request.generation).await?;
+    let recovery = if request.ignore_state {
+        RecoverySearch::Found(GenerationRecovery::NoCheckpoint)
+    } else {
+        find_recovery_checkpoint(store, &paths, request.generation).await?
+    };
     let base_checkpoint_ref = match &recovery {
         RecoverySearch::Found(recovery) => match recovery {
             GenerationRecovery::NoCheckpoint => None,

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

@@ -160,6 +160,9 @@ impl WorkerJobController {
                 job_id: worker_context.job_id.clone(),
                 generation: Generation(worker_context.generation),
                 updated_at: SystemTime::now(),
+                // The controller passes no parent checkpoint when this generation should start
+                // without state (or when there is no state available).
+                ignore_state: parent_ref.is_none(),
             },
             false,
         )

webui/src/routes/pipelines/PipelineDetails.tsx

@@ -138,7 +138,7 @@ export function PipelineDetails() {
 
   async function updateJobState(stop: StopType) {
     console.log(`Setting pipeline stop_mode=${stop}`);
-    updatePipeline({ stop });
+    await updatePipeline({ stop });
   }
 
   async function updateJobParallelism(parallelism: number) {
@@ -352,23 +352,29 @@ export function PipelineDetails() {
   let actionButton = <></>;
   if (pipeline) {
     editPipelineButton = <Button onClick={onConfigModalOpen}>Edit</Button>;
-    if (job.state == 'Failed') {
+    const isStopped = job.state === 'Stopped' && pipeline.action != null;
+    if (job.state === 'Failed' || isStopped) {
+      const actionText = isStopped ? 'Start' : 'Restart';
       actionButton = (
         <Popover trigger="hover" placement="bottom-start">
           <PopoverTrigger>
             <Button
               onClick={async () => {
-                await restartPipeline(false);
+                if (isStopped) {
+                  await updateJobState(pipeline.action!);
+                } else {
+                  await restartPipeline(false);
+                }
               }}
             >
-              Restart
+              {actionText}
             </Button>
           </PopoverTrigger>
           <PopoverContent width="auto">
             <PopoverArrow />
             <PopoverBody p={4}>
               <Button size="sm" colorScheme="red" onClick={onRestartWithoutStateModalOpen}>
-                Restart Without State
+                {actionText} Without State
               </Button>
             </PopoverBody>
           </PopoverContent>
@@ -389,6 +395,10 @@ export function PipelineDetails() {
     }
   }
 
+  const startingWithoutState = job.state === 'Stopped';
+  const withoutStateAction = startingWithoutState ? 'Start' : 'Restart';
+  const withoutStateActionPresentParticiple = startingWithoutState ? 'Starting' : 'Restarting';
+
   const headerArea = (
     <Flex>
       <Box p={5}>
@@ -422,13 +432,13 @@ export function PipelineDetails() {
         <AlertDialogOverlay>
           <AlertDialogContent>
             <AlertDialogHeader fontSize="lg" fontWeight="bold">
-              Restart Without State
+              {withoutStateAction} Without State
             </AlertDialogHeader>
 
             <AlertDialogBody>
               <Text>
-                Restarting without state could lead to data loss, duplication, and incorrect
-                results. Are you sure you want to continue?
+                {withoutStateActionPresentParticiple} without state could lead to data loss,
+                duplication, and incorrect results. Are you sure you want to continue?
               </Text>
             </AlertDialogBody>
 
@@ -444,7 +454,7 @@ export function PipelineDetails() {
                 }}
                 ml={3}
               >
-                Restart Without State
+                {withoutStateAction} Without State
               </Button>
             </AlertDialogFooter>
           </AlertDialogContent>