Skip to content
This repository was archived by the owner on Oct 19, 2026. It is now read-only.

Commit 48fc759

Browse files
topshihunCopilot
andauthored
feat: Auto retry judge and build (#1768)
* feat: auto retry time Signed-off-by: Wangyan <topshihun@foxmail.com> * Update orion-server/src/auto_retry.rs Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Signed-off-by: TOP <89700178+topshihun@users.noreply.github.com> * fix: clippy and other advices Signed-off-by: Wangyan <topshihun@foxmail.com> --------- Signed-off-by: Wangyan <topshihun@foxmail.com> Signed-off-by: TOP <89700178+topshihun@users.noreply.github.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
1 parent 57ceddb commit 48fc759

11 files changed

Lines changed: 406 additions & 73 deletions

File tree

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
use sea_orm_migration::prelude::*;
2+
3+
#[derive(DeriveMigrationName)]
4+
pub struct Migration;
5+
6+
#[async_trait::async_trait]
7+
impl MigrationTrait for Migration {
8+
async fn up(&self, manager: &SchemaManager) -> Result<(), DbErr> {
9+
manager
10+
.alter_table(
11+
Table::alter()
12+
.table(Builds::Table)
13+
.add_column(
14+
ColumnDef::new(Alias::new("retry_count"))
15+
.integer()
16+
.not_null()
17+
.default(0),
18+
)
19+
.to_owned(),
20+
)
21+
.await?;
22+
Ok(())
23+
}
24+
25+
async fn down(&self, manager: &SchemaManager) -> Result<(), DbErr> {
26+
manager
27+
.alter_table(
28+
Table::alter()
29+
.table(Builds::Table)
30+
.drop_column(Alias::new("retry_count"))
31+
.to_owned(),
32+
)
33+
.await?;
34+
Ok(())
35+
}
36+
}
37+
38+
#[derive(DeriveIden)]
39+
enum Builds {
40+
Table,
41+
}

‎jupiter/src/migration/mod.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,7 @@ mod m20251125_135032_add_draft_conv_type;
7474
mod m20251128_000001_create_buck_session;
7575
mod m20251203_013745_add_dynamic_sidebar;
7676
mod m20251210_113942_remove_unique_constraint_from_order_index;
77+
mod m20260106_070511_add_retry_time;
7778
mod m20260106_070515_remove_relay_mq_lfs_raw_table;
7879

7980
/// Creates a primary key column definition with big integer type.
@@ -137,6 +138,7 @@ impl MigratorTrait for Migrator {
137138
Box::new(m20251128_000001_create_buck_session::Migration),
138139
Box::new(m20251203_013745_add_dynamic_sidebar::Migration),
139140
Box::new(m20251210_113942_remove_unique_constraint_from_order_index::Migration),
141+
Box::new(m20260106_070511_add_retry_time::Migration),
140142
Box::new(m20260106_070515_remove_relay_mq_lfs_raw_table::Migration),
141143
]
142144
}

‎orion-server/Cargo.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ anyhow = { workspace = true }
3636
async-trait = { workspace = true }
3737
aws-sdk-s3 = { workspace = true }
3838
aws-config = { workspace = true }
39+
once_cell = { workspace = true }
3940

4041
[dev-dependencies]
4142
tempfile = "3"

‎orion-server/db/schema.sql‎

Lines changed: 0 additions & 37 deletions
This file was deleted.

‎orion-server/src/api.rs‎

Lines changed: 138 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
use crate::auto_retry::AutoRetryJudger;
12
use crate::common::model::{CommonPage, PageParams};
23
use crate::log::log_service::{LogEvent, LogService};
34
use crate::model::{builds, tasks};
@@ -38,6 +39,8 @@ use tokio_stream::wrappers::IntervalStream;
3839
use utoipa::ToSchema;
3940
use uuid::Uuid;
4041

42+
const RETRY_COUNT_MAX: u32 = 3;
43+
4144
#[derive(Debug, Serialize, Default, ToSchema)]
4245
struct BuildResult {
4346
build_id: String,
@@ -70,6 +73,8 @@ pub struct TaskRequest {
7073
pub repo: String,
7174
pub cl_link: String,
7275
pub cl: i64,
76+
pub task_name: Option<String>,
77+
pub template: Option<Value>,
7378
pub builds: Vec<scheduler::BuildRequest>,
7479
}
7580

@@ -298,8 +303,8 @@ pub async fn task_handler(
298303
if let Err(err) = tasks::Model::insert_task(
299304
task_id,
300305
req.cl,
301-
None,
302-
None,
306+
req.task_name.clone(),
307+
req.template.clone(),
303308
chrono::Utc::now().into(),
304309
&state.conn,
305310
)
@@ -418,6 +423,8 @@ async fn handle_immediate_task_dispatch(
418423
start_at: chrono::Utc::now(),
419424
cl: cl.to_string(),
420425
_worker_id: chosen_id.clone(),
426+
auto_retry_judger: AutoRetryJudger::new(),
427+
retry_count: 0,
421428
};
422429

423430
// Use the model's insert_build method for direct insertion
@@ -693,6 +700,11 @@ async fn process_message(
693700
} else {
694701
tracing::warn!("Received output for unknown task: {}", id);
695702
}
703+
704+
// Judge auto retry by output
705+
if let Some(mut build_info) = state.scheduler.active_builds.get_mut(&id) {
706+
build_info.auto_retry_judger.judge_by_output(&output);
707+
}
696708
}
697709
WSMessage::BuildComplete {
698710
id,
@@ -705,44 +717,134 @@ async fn process_message(
705717
"Build {id} completed by worker {current_worker_id} with exit code: {exit_code:?}"
706718
);
707719

708-
// Send final log event
720+
// Get build information
709721
if let Some(build_info) = state.scheduler.active_builds.get(&id) {
710-
let log_event = LogEvent {
711-
task_id: build_info.task_id.clone(),
712-
repo_name: LogService::last_segment(&build_info.repo.clone())
713-
.to_string(),
714-
build_id: build_info.build_id.clone(),
715-
line: String::from(""), // Empty line to signal end
716-
is_end: true, // This is the end of the log
717-
};
718-
state.log_service.publish(log_event);
719-
}
722+
let mut build_info = build_info.clone();
720723

721-
// Remove from active builds and update database
722-
state.scheduler.active_builds.remove(&id);
723-
let _ = builds::Entity::update_many()
724-
.set(builds::ActiveModel {
725-
exit_code: Set(exit_code),
726-
end_at: Set(Some(
727-
Utc::now().with_timezone(&FixedOffset::east_opt(0).unwrap()),
728-
)),
729-
..Default::default()
730-
})
731-
.filter(builds::Column::Id.eq(id.parse::<uuid::Uuid>().unwrap()))
732-
.exec(&state.conn)
733-
.await;
734-
735-
// Mark the worker as idle or error depending on whether the task succeeds.
736-
if let Some(mut worker) = state.scheduler.workers.get_mut(current_worker_id) {
737-
worker.status = if success {
738-
WorkerStatus::Idle
724+
// Judge auto retry by exit code
725+
build_info
726+
.auto_retry_judger
727+
.judge_by_exit_code(exit_code.unwrap_or(0));
728+
729+
let (can_auto_retry, mut retry_count) = (
730+
build_info.auto_retry_judger.get_can_auto_retry(),
731+
build_info.retry_count,
732+
);
733+
734+
// Reset auto retry judger
735+
build_info.auto_retry_judger = AutoRetryJudger::new();
736+
737+
if can_auto_retry && retry_count < RETRY_COUNT_MAX {
738+
// Restart this build for retry and add retry time
739+
retry_count += 1;
740+
741+
// Update build information
742+
build_info.retry_count = retry_count;
743+
state
744+
.scheduler
745+
.active_builds
746+
.alter(&id, |_, _| build_info.clone());
747+
748+
// Update database
749+
let _ = builds::Entity::update_many()
750+
.set(builds::ActiveModel {
751+
retry_count: Set(retry_count),
752+
..Default::default()
753+
})
754+
.filter(builds::Column::Id.eq(id.parse::<uuid::Uuid>().unwrap()))
755+
.exec(&state.conn)
756+
.await;
757+
758+
// Send Task to the same worker
759+
let msg = WSMessage::Task {
760+
id: build_info.build_id.clone(),
761+
repo: build_info.repo,
762+
changes: build_info.changes,
763+
cl_link: build_info.cl,
764+
};
765+
let worker_id = build_info._worker_id;
766+
if let Some(worker) = state.scheduler.workers.get_mut(&worker_id)
767+
&& worker.sender.send(msg).is_ok()
768+
{
769+
tracing::info!(
770+
"Retry build: {}, worker: {}",
771+
build_info.build_id,
772+
worker_id
773+
);
774+
} else {
775+
tracing::error!("Retry build send to worker Failed");
776+
777+
// If retry dispatch fails, treat this build as finished and clean up.
778+
// Remove from active builds
779+
state.scheduler.active_builds.remove(&id);
780+
781+
// Update database with final state
782+
let _ = builds::Entity::update_many()
783+
.set(builds::ActiveModel {
784+
exit_code: Set(exit_code),
785+
end_at: Set(Some(
786+
Utc::now()
787+
.with_timezone(&FixedOffset::east_opt(0).unwrap()),
788+
)),
789+
retry_count: Set(retry_count),
790+
..Default::default()
791+
})
792+
.filter(
793+
builds::Column::Id.eq(id.parse::<uuid::Uuid>().unwrap()),
794+
)
795+
.exec(&state.conn)
796+
.await;
797+
798+
// Notify scheduler to process queued tasks
799+
state.scheduler.notify_task_available();
800+
}
739801
} else {
740-
WorkerStatus::Error(message)
741-
};
742-
}
802+
// Send final log event
803+
let log_event = LogEvent {
804+
task_id: build_info.task_id.clone(),
805+
repo_name: LogService::last_segment(&build_info.repo.clone())
806+
.to_string(),
807+
build_id: build_info.build_id.clone(),
808+
line: String::from(""),
809+
is_end: true,
810+
};
811+
state.log_service.publish(log_event);
812+
813+
// Remove from active builds
814+
state.scheduler.active_builds.remove(&id);
815+
816+
// Update database
817+
let _ = builds::Entity::update_many()
818+
.set(builds::ActiveModel {
819+
exit_code: Set(exit_code),
820+
end_at: Set(Some(
821+
Utc::now()
822+
.with_timezone(&FixedOffset::east_opt(0).unwrap()),
823+
)),
824+
retry_count: Set(retry_count),
825+
..Default::default()
826+
})
827+
.filter(builds::Column::Id.eq(id.parse::<uuid::Uuid>().unwrap()))
828+
.exec(&state.conn)
829+
.await;
830+
831+
// Mark the worker as idle or error depending on whether the task succeeds.
832+
if let Some(mut worker) =
833+
state.scheduler.workers.get_mut(current_worker_id)
834+
{
835+
worker.status = if success {
836+
WorkerStatus::Idle
837+
} else {
838+
WorkerStatus::Error(message)
839+
};
840+
}
743841

744-
// After worker becomes idle, notify to process queued tasks
745-
state.scheduler.notify_task_available();
842+
// After worker becomes idle, notify to process queued tasks
843+
state.scheduler.notify_task_available();
844+
}
845+
} else {
846+
tracing::error!("Not found build: {id}");
847+
}
746848
}
747849
WSMessage::TaskPhaseUpdate { id, phase } => {
748850
tracing::info!(

0 commit comments

Comments
 (0)