1+ use crate :: auto_retry:: AutoRetryJudger ;
12use crate :: common:: model:: { CommonPage , PageParams } ;
23use crate :: log:: log_service:: { LogEvent , LogService } ;
34use crate :: model:: { builds, tasks} ;
@@ -38,6 +39,8 @@ use tokio_stream::wrappers::IntervalStream;
3839use utoipa:: ToSchema ;
3940use uuid:: Uuid ;
4041
42+ const RETRY_COUNT_MAX : u32 = 3 ;
43+
4144#[ derive( Debug , Serialize , Default , ToSchema ) ]
4245struct 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