Skip to content

Commit fa575e4

Browse files
mwyldemwylde
authored
Add support for starting/restarting without state in leader mode (#1122)
1 parent f6afb83 commit fa575e4

10 files changed

Lines changed: 413 additions & 105 deletions

File tree

‎crates/arroyo-api/queries/api_queries.sql‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -183,7 +183,6 @@ UPDATE job_configs
183183
SET
184184
updated_at = :updated_at,
185185
updated_by = :updated_by,
186-
stop = 'none',
187186
restart_nonce = restart_nonce + 1,
188187
restart_mode = :mode,
189188
ignore_state_before_epoch = :ignore_state_before_epoch

‎crates/arroyo-api/src/jobs.rs‎

Lines changed: 383 additions & 1 deletion
Large diffs are not rendered by default.

‎crates/arroyo-api/src/pipelines.rs‎

Lines changed: 21 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -869,43 +869,39 @@ pub async fn restart_pipeline(
869869
) -> Result<Json<Pipeline>, ErrorResp> {
870870
let auth_data = authenticate(&state.database, bearer_auth).await?;
871871

872+
if req.ignore_state.unwrap_or(false)
873+
&& matches!(config().job_controller, JobControllerMode::Worker)
874+
{
875+
jobs::replace_job_without_state(&state.database, &id, &auth_data).await?;
876+
877+
let db = state.database.client().await?;
878+
let pipeline = query_pipeline_by_pub_id(&id, &db, &auth_data).await?;
879+
return Ok(Json(pipeline));
880+
}
881+
872882
let db = state.database.client().await?;
873883

874-
let job = api_queries::fetch_get_pipeline_jobs(&db, &auth_data.organization_id, &id)
884+
let job_id = api_queries::fetch_get_pipeline_jobs(&db, &auth_data.organization_id, &id)
875885
.await?
876886
.into_iter()
877887
.next()
878-
.ok_or_else(|| bad_request("No jobs for pipeline"))?;
888+
.ok_or_else(|| bad_request("No jobs for pipeline"))?
889+
.id;
879890

880891
let mode = if req.force == Some(true) {
881892
RestartMode::force
882893
} else {
883894
RestartMode::safe
884895
};
885896

897+
// If user wants to ignore state, query max checkpoint epoch and compute threshold
886898
let ignore_before_epoch = if req.ignore_state.unwrap_or(false) {
887-
match config().job_controller {
888-
JobControllerMode::Controller => {
889-
// Controller mode uses this as an epoch threshold.
890-
api_queries::fetch_max_checkpoint_epoch(&db, &job.id, &auth_data.organization_id)
891-
.await?
892-
.into_iter()
893-
.next()
894-
.and_then(|r| r.max_epoch)
895-
.map(|max_epoch| max_epoch + 1)
896-
}
897-
JobControllerMode::Worker => {
898-
// Leader mode uses this as the generation that should start without state.
899-
Some(
900-
job.run_id
901-
.unwrap_or(0)
902-
.max(0)
903-
.checked_add(1)
904-
.and_then(|generation| generation.try_into().ok())
905-
.ok_or_else(|| bad_request("Job generation is too large to restart"))?,
906-
)
907-
}
908-
}
899+
api_queries::fetch_max_checkpoint_epoch(&db, &job_id, &auth_data.organization_id)
900+
.await?
901+
.into_iter()
902+
.next()
903+
.and_then(|r| r.max_epoch)
904+
.map(|max_epoch| max_epoch + 1)
909905
} else {
910906
None
911907
};
@@ -916,7 +912,7 @@ pub async fn restart_pipeline(
916912
&auth_data.user_id,
917913
&mode,
918914
&ignore_before_epoch,
919-
&job.id,
915+
&job_id,
920916
&auth_data.organization_id,
921917
)
922918
.await?;

‎crates/arroyo-api/src/rest_utils.rs‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -168,6 +168,13 @@ pub(crate) fn bad_request(message: impl Into<String>) -> ErrorResp {
168168
}
169169
}
170170

171+
pub(crate) fn conflict(message: impl Into<String>) -> ErrorResp {
172+
ErrorResp {
173+
status_code: StatusCode::CONFLICT,
174+
message: message.into(),
175+
}
176+
}
177+
171178
pub(crate) fn service_unavailable(object: &str) -> ErrorResp {
172179
ErrorResp {
173180
status_code: StatusCode::SERVICE_UNAVAILABLE,

‎crates/arroyo-controller/src/lib.rs‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -111,8 +111,6 @@ pub struct JobConfig {
111111
parallelism_overrides: HashMap<u32, usize>,
112112
restart_nonce: i32,
113113
restart_mode: RestartMode,
114-
/// Minimum checkpoint epoch in controller mode; generation to start without state in leader
115-
/// mode.
116114
ignore_state_before_epoch: Option<i32>,
117115
/// Per-job environment variables forwarded to workers at scheduling time.
118116
env_vars: serde_json::Value,

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

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -472,7 +472,6 @@ async fn get_and_register_checkpoint_info_leader<'a>(
472472
job_id: JobId(ctx.config.id.clone()),
473473
generation: Generation(ctx.status.generation),
474474
updated_at: SystemTime::now(),
475-
ignore_state,
476475
},
477476
true,
478477
)

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

Lines changed: 0 additions & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -759,7 +759,6 @@ mod tests {
759759
job_id: JobId::new("J"),
760760
generation: Generation(1),
761761
updated_at: from_micros(123),
762-
ignore_state: false,
763762
},
764763
false,
765764
)
@@ -815,7 +814,6 @@ mod tests {
815814
job_id: JobId::new("J"),
816815
generation: Generation(2),
817816
updated_at: from_micros(456),
818-
ignore_state: false,
819817
},
820818
false,
821819
)
@@ -847,62 +845,6 @@ mod tests {
847845
assert_eq!(written_manifest, expected_manifest);
848846
}
849847

850-
#[tokio::test]
851-
async fn initialize_generation_can_ignore_previous_checkpoint() {
852-
let store = MemoryProtocolStore::default();
853-
let paths = ProtocolPaths::new(PipelineId::new("P"), JobId::new("J"));
854-
write_current_generation(&store, &paths, Generation(2)).await;
855-
856-
let checkpoint_ref = paths.checkpoint_manifest(Generation(1), Epoch(1));
857-
let checkpoint = checkpoint_for_generation(Generation(1), 1, None, false);
858-
write_canonical_checkpoint(&store, &paths, &checkpoint_ref, &checkpoint).await;
859-
let previous_manifest =
860-
generation_manifest_for_generation(Generation(1), None, Some(checkpoint_ref));
861-
put_json(
862-
&store,
863-
&paths.generation_manifest(Generation(1)),
864-
&previous_manifest,
865-
)
866-
.await
867-
.unwrap();
868-
869-
let initialization = initialize_generation(
870-
&store,
871-
InitializeGenerationRequest {
872-
pipeline_id: PipelineId::new("P"),
873-
job_id: JobId::new("J"),
874-
generation: Generation(2),
875-
updated_at: from_micros(456),
876-
ignore_state: true,
877-
},
878-
false,
879-
)
880-
.await
881-
.unwrap();
882-
883-
let expected_manifest = GenerationManifest::new(
884-
PipelineId::new("P"),
885-
JobId::new("J"),
886-
Generation(2),
887-
None,
888-
456,
889-
);
890-
assert_eq!(
891-
initialization,
892-
GenerationInitialization::Initialized {
893-
generation_manifest: expected_manifest.clone(),
894-
recovery: GenerationRecovery::NoCheckpoint,
895-
}
896-
);
897-
898-
let written_manifest: GenerationManifest =
899-
read_json(&store, &paths.generation_manifest(Generation(2)))
900-
.await
901-
.unwrap()
902-
.expect("new generation manifest should be written");
903-
assert_eq!(written_manifest, expected_manifest);
904-
}
905-
906848
#[tokio::test]
907849
async fn initialize_generation_restores_previous_checkpoint_requiring_commit_replay() {
908850
let store = MemoryProtocolStore::default();
@@ -929,7 +871,6 @@ mod tests {
929871
job_id: JobId::new("J"),
930872
generation: Generation(2),
931873
updated_at: from_micros(456),
932-
ignore_state: false,
933874
},
934875
false,
935876
)
@@ -978,7 +919,6 @@ mod tests {
978919
job_id: JobId::new("J"),
979920
generation: Generation(3),
980921
updated_at: from_micros(789),
981-
ignore_state: false,
982922
},
983923
false,
984924
)
@@ -1026,7 +966,6 @@ mod tests {
1026966
job_id: JobId::new("J"),
1027967
generation: Generation(2),
1028968
updated_at: from_micros(456),
1029-
ignore_state: false,
1030969
},
1031970
false,
1032971
)
@@ -1083,7 +1022,6 @@ mod tests {
10831022
job_id: JobId::new("J"),
10841023
generation: Generation(3),
10851024
updated_at: from_micros(456),
1086-
ignore_state: false,
10871025
},
10881026
false,
10891027
)
@@ -1149,7 +1087,6 @@ mod tests {
11491087
job_id: JobId::new("J"),
11501088
generation: Generation(3),
11511089
updated_at: from_micros(456),
1152-
ignore_state: false,
11531090
},
11541091
false,
11551092
)
@@ -1187,7 +1124,6 @@ mod tests {
11871124
job_id: JobId::new("J"),
11881125
generation: Generation(2),
11891126
updated_at: from_micros(456),
1190-
ignore_state: false,
11911127
},
11921128
false,
11931129
)

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

Lines changed: 1 addition & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -71,8 +71,6 @@ pub struct InitializeGenerationRequest {
7171
pub job_id: JobId,
7272
pub generation: Generation,
7373
pub updated_at: SystemTime,
74-
/// Start this generation without restoring a checkpoint from an earlier generation.
75-
pub ignore_state: bool,
7674
}
7775

7876
/// Checkpoint, if any, that a newly initialized generation should restore from.
@@ -261,11 +259,7 @@ where
261259
});
262260
}
263261

264-
let recovery = if request.ignore_state {
265-
RecoverySearch::Found(GenerationRecovery::NoCheckpoint)
266-
} else {
267-
find_recovery_checkpoint(store, &paths, request.generation).await?
268-
};
262+
let recovery = find_recovery_checkpoint(store, &paths, request.generation).await?;
269263
let base_checkpoint_ref = match &recovery {
270264
RecoverySearch::Found(recovery) => match recovery {
271265
GenerationRecovery::NoCheckpoint => None,

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

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -160,9 +160,6 @@ impl WorkerJobController {
160160
job_id: worker_context.job_id.clone(),
161161
generation: Generation(worker_context.generation),
162162
updated_at: SystemTime::now(),
163-
// The controller passes no parent checkpoint when this generation should start
164-
// without state (or when there is no state available).
165-
ignore_state: parent_ref.is_none(),
166163
},
167164
false,
168165
)

‎webui/src/lib/data_fetching.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -524,7 +524,7 @@ export const usePipeline = (pipelineId?: string, refresh: boolean = false) => {
524524
params: { path: { id: pipelineId } },
525525
body: { ignore_state: ignoreState ?? false },
526526
});
527-
await mutate();
527+
await Promise.all([mutate(), globalMutate(pipelineJobsKey(pipelineId))]);
528528
};
529529

530530
const deletePipeline = async () => {

0 commit comments

Comments
 (0)