Support restarting without state in leader mode

Author: mwylde

crates/arroyo-api/queries/api_queries.sql

@@ -196,6 +196,32 @@ VALUES (:id, :organization_id, :pipeline_name, :created_by, :pipeline_id, :check
 --! create_job_status
 INSERT INTO job_statuses (pub_id, id, organization_id) VALUES (:pub_id, :id, :organization_id);
 
+--! clone_job_for_restart
+INSERT INTO job_configs
+(id, organization_id, pipeline_name, created_by, updated_by, updated_at, ttl_micros, stop,
+ pipeline_id, parallelism_overrides, checkpoint_interval_micros, env_vars, scheduler_config)
+SELECT :new_job_id, organization_id, pipeline_name, created_by, :updated_by, CURRENT_TIMESTAMP,
+       ttl_micros, 'none', pipeline_id, parallelism_overrides, checkpoint_interval_micros,
+       env_vars, scheduler_config
+FROM job_configs
+WHERE id = :old_job_id AND organization_id = :organization_id;
+
+--! delete_job_checkpoints
+DELETE FROM checkpoints
+WHERE job_id = :job_id AND organization_id = :organization_id;
+
+--! delete_job_log_messages
+DELETE FROM job_log_messages
+WHERE job_id = :job_id;
+
+--! delete_job_status
+DELETE FROM job_statuses
+WHERE id = :job_id AND organization_id = :organization_id;
+
+--! delete_job_config
+DELETE FROM job_configs
+WHERE id = :job_id AND organization_id = :organization_id;
+
 --! get_jobs: (start_time?, finish_time?, state?, tasks?, textual_repr?, failure_message?, run_id?, udfs)
 SELECT job_configs.id as id, pipeline_name, stop, textual_repr, start_time, finish_time, state, tasks, pipeline_id, failure_message, run_id, udfs
 FROM job_configs

crates/arroyo-api/src/jobs.rs

@@ -29,7 +29,8 @@ use crate::pipelines::{query_job_by_pub_id, query_pipeline_by_pub_id};
 use crate::rest::AppState;
 use crate::rest_utils::{
     BearerAuth, ErrorResp, PipelineJobCheckpointPath, PipelineJobPath, authenticate, bad_request,
-    log_and_map, not_found, paginate_results, validate_pagination_params,
+    conflict, internal_server_error, log_and_map, not_found, paginate_results,
+    validate_pagination_params,
 };
 use crate::types::public::LogLevel;
 use crate::{AuthData, queries::api_queries, to_micros, types::public};
@@ -194,6 +195,76 @@ pub(crate) async fn create_job(
     Ok(job_id)
 }
 
+fn replaceable_job(jobs: Vec<DbPipelineJob>) -> Result<String, ErrorResp> {
+    let count = jobs.len();
+    let Some(job) = jobs.into_iter().next() else {
+        return Err(not_found("Job for pipeline"));
+    };
+
+    if count != 1 {
+        return Err(internal_server_error(format!(
+            "expected one job for pipeline, found {count}"
+        )));
+    }
+
+    let state = job.state.unwrap_or_else(|| "Created".to_string());
+    if state != "Stopped" && state != "Failed" {
+        return Err(conflict(format!(
+            "cannot restart job {} without state while it is in state {state}; stop the job first",
+            job.id
+        )));
+    }
+
+    Ok(job.id)
+}
+
+/// Replaces a pipeline's single terminal job with a fresh job.
+pub(crate) async fn replace_job_without_state(
+    db: &DatabaseSource,
+    pipeline_pub_id: &str,
+    auth: &AuthData,
+) -> Result<String, ErrorResp> {
+    let new_job_id = generate_id(IdTypes::JobConfig);
+    let new_status_id = generate_id(IdTypes::JobStatus);
+    let database = db.client().await?;
+    let jobs =
+        api_queries::fetch_get_pipeline_jobs(&database, &auth.organization_id, &pipeline_pub_id)
+            .await?;
+    let old_job_id = replaceable_job(jobs)?;
+
+    // TODO: this really should all be within a transaction, however we can't use transactions
+    //  today due to our sqlite/postgres compatibility layer
+    let inserted = api_queries::execute_clone_job_for_restart(
+        &database,
+        &new_job_id,
+        &auth.user_id,
+        &old_job_id,
+        &auth.organization_id,
+    )
+    .await?;
+    if inserted != 1 {
+        return Err(internal_server_error("failed to clone job configuration"));
+    }
+
+    api_queries::execute_create_job_status(
+        &database,
+        &new_status_id,
+        &new_job_id,
+        &auth.organization_id,
+    )
+    .await?;
+
+    // Delete children explicitly because the SQLite checkpoints table does not have the same
+    // foreign-key cascade as PostgreSQL.
+    api_queries::execute_delete_job_checkpoints(&database, &old_job_id, &auth.organization_id)
+        .await?;
+    api_queries::execute_delete_job_log_messages(&database, &old_job_id).await?;
+    api_queries::execute_delete_job_status(&database, &old_job_id, &auth.organization_id).await?;
+    api_queries::execute_delete_job_config(&database, &old_job_id, &auth.organization_id).await?;
+
+    Ok(new_job_id)
+}
+
 pub(crate) fn get_action(state: &str, running_desired: &bool) -> (String, Option<StopType>, bool) {
     enum Progress {
         InProgress,

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,6 +868,17 @@ 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?;
+
+    if req.ignore_state.unwrap_or(false)
+        && matches!(config().job_controller, JobControllerMode::Worker)
+    {
+        jobs::replace_job_without_state(&state.database, &id, &auth_data).await?;
+
+        let db = state.database.client().await?;
+        let pipeline = query_pipeline_by_pub_id(&id, &db, &auth_data).await?;
+        return Ok(Json(pipeline));
+    }
+
     let db = state.database.client().await?;
 
     let job_id = api_queries::fetch_get_pipeline_jobs(&db, &auth_data.organization_id, &id)

crates/arroyo-api/src/rest_utils.rs

@@ -168,6 +168,13 @@ pub(crate) fn bad_request(message: impl Into<String>) -> ErrorResp {
     }
 }
 
+pub(crate) fn conflict(message: impl Into<String>) -> ErrorResp {
+    ErrorResp {
+        status_code: StatusCode::CONFLICT,
+        message: message.into(),
+    }
+}
+
 pub(crate) fn service_unavailable(object: &str) -> ErrorResp {
     ErrorResp {
         status_code: StatusCode::SERVICE_UNAVAILABLE,

webui/src/lib/data_fetching.ts

@@ -524,7 +524,7 @@ export const usePipeline = (pipelineId?: string, refresh: boolean = false) => {
       params: { path: { id: pipelineId } },
       body: { ignore_state: ignoreState ?? false },
     });
-    await mutate();
+    await Promise.all([mutate(), globalMutate(pipelineJobsKey(pipelineId))]);
   };
 
   const deletePipeline = async () => {

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>