refactor: add reconnect as a backend field, always create RecoveryManager

Author: hhamalai

agent/runner.go

@@ -94,18 +94,28 @@ func (r *Runner) Run(runnerCtx, shutdownCtx context.Context) error {
 	workflowCtx, cancelWorkflowCtx := context.WithCancelCause(workflowCtx)
 	defer cancelWorkflowCtx(nil)
 
-	// recoveryManager is declared here so the cancel listener can mark it as canceled.
-	// It will be initialized later after workflow state is set up.
-	var recoveryManager *pipeline.RecoveryManager
-
-	// Add sigterm support for internal context.
-	// Required to be able to terminate the running workflow by external signals.
+	// Handle SIGTERM (k8s, docker, system shutdown)
 	workflowCtx = utils.WithContextSigtermCallback(workflowCtx, func() {
 		logger.Error().Msg("received sigterm termination signal")
 		// WithContextSigtermCallback would cancel the context too, but  we want our own custom error
 		cancelWorkflowCtx(pipeline.ErrCancel)
 	})
 
+	state := rpc.WorkflowState{
+		Started: time.Now().Unix(),
+	}
+	if err := r.client.Init(runnerCtx, workflow.ID, state); err != nil {
+		logger.Error().Err(err).Msg("workflow initialization failed")
+		// TODO: should we return here?
+	}
+
+	// Initialize recovery manager before launching goroutines that reference it
+	recoveryManager := pipeline.NewRecoveryManager(r.client, workflow.ID, true)
+	if err := recoveryManager.InitRecoveryState(runnerCtx, workflow.Config, int64(timeout.Seconds())); err != nil {
+		logger.Warn().Err(err).Msg("failed to initialize recovery state, continuing without recovery")
+		recoveryManager = pipeline.NewRecoveryManager(r.client, workflow.ID, false)
+	}
+
 	// Listen for remote cancel events (UI / API).
 	// When canceled, we MUST cancel the workflow context
 	// so that workflow execution stop immediately.
@@ -118,9 +128,7 @@ func (r *Runner) Run(runnerCtx, shutdownCtx context.Context) error {
 		} else {
 			if canceled {
 				logger.Debug().Msg("server side cancel signal received")
-				if recoveryManager != nil {
-					recoveryManager.SetCanceled()
-				}
+				recoveryManager.SetCanceled()
 				cancelWorkflowCtx(pipeline.ErrCancel)
 			}
 			logger.Debug().Msg("cancel listener exited normally")
@@ -144,25 +152,6 @@ func (r *Runner) Run(runnerCtx, shutdownCtx context.Context) error {
 		}
 	}()
 
-	state := rpc.WorkflowState{
-		Started: time.Now().Unix(),
-	}
-
-	if err := r.client.Init(runnerCtx, workflow.ID, state); err != nil {
-		logger.Error().Err(err).Msg("signaling workflow initialization to server failed")
-		// We have an error, maybe the server is currently unreachable or other server-side errors occurred.
-		// So let's clean up and end this not yet started workflow run.
-		cancelWorkflowCtx(err)
-		return err
-	}
-
-	// Initialize recovery manager; if not enabled on server, it will be a no-op
-	recoveryManager = pipeline.NewRecoveryManager(r.client, workflow.ID, true)
-	if err := recoveryManager.InitRecoveryState(runnerCtx, workflow.Config, int64(timeout.Seconds())); err != nil {
-		logger.Warn().Err(err).Msg("failed to initialize recovery state, continuing without recovery")
-		recoveryManager = pipeline.NewRecoveryManager(nil, workflow.ID, false)
-	}
-
 	var uploads sync.WaitGroup
 
 	// Run pipeline
@@ -204,7 +193,7 @@ func (r *Runner) Run(runnerCtx, shutdownCtx context.Context) error {
 
 	// If workflow is recoverable (context canceled, recovery enabled, not user cancel),
 	// skip marking as done. The workflow will be picked up by a new agent after restart.
-	if recoveryManager != nil && recoveryManager.IsRecoverable(runnerCtx) {
+	if recoveryManager.IsRecoverable(runnerCtx) {
 		logger.Info().Msg("workflow is recoverable, not marking as done")
 		return nil
 	}

pipeline/backend/dummy/dummy.go

@@ -216,6 +216,10 @@ func (e *dummy) DestroyStep(_ context.Context, step *backend.Step, taskUUID stri
 	return nil
 }
 
+func (e *dummy) Reconnect(_ context.Context, _ *backend.Step, _ string) error {
+	return fmt.Errorf("reconnect not supported")
+}
+
 func (e *dummy) DestroyWorkflow(_ context.Context, _ *backend.Config, taskUUID string) error {
 	log.Trace().Str("taskUUID", taskUUID).Msgf("delete workflow environment")
 

pipeline/backend/local/local.go

@@ -256,6 +256,10 @@ func (e *local) DestroyStep(_ context.Context, step *types.Step, taskUUID string
 	return nil
 }
 
+func (e *local) Reconnect(_ context.Context, _ *types.Step, _ string) error {
+	return fmt.Errorf("reconnect not supported")
+}
+
 func (e *local) DestroyWorkflow(_ context.Context, _ *types.Config, taskUUID string) error {
 	log.Trace().Str("taskUUID", taskUUID).Msg("delete workflow environment")
 

pipeline/backend/types/backend.go

@@ -161,18 +161,15 @@ type Backend interface {
 	// This function may be called concurrently for different workflows
 	// and must be thread-safe.
 	DestroyWorkflow(ctx context.Context, conf *Config, taskUUID string) error
+
+	// Reconnect attempts to reconnect to a running step after agent restart.
+	// Returns nil if reconnection is possible, error otherwise.
+	// After successful reconnect, TailStep and WaitStep can be used normally.
+	// Backends that do not support reconnection should return an error.
+	Reconnect(ctx context.Context, step *Step, taskUUID string) error
 }
 
 // BackendInfo represents the reported information of a loaded backend.
 type BackendInfo struct {
 	Platform string
 }
-
-// Reconnector is an optional interface that backends can implement to support
-// reconnecting to running steps after agent restart.
-type Reconnector interface {
-	// Reconnect attempts to reconnect to a running step.
-	// Returns nil if reconnection is possible, error otherwise.
-	// After successful reconnect, TailStep and WaitStep can be used normally.
-	Reconnect(ctx context.Context, step *Step, taskUUID string) error
-}

pipeline/pipeline.go

@@ -256,16 +256,12 @@ func (r *Runtime) execAll(runnerCtx context.Context, steps []*backend.Step) <-ch
 					return nil
 				} else if r.recoveryManager.ShouldReconnect(recoveryState) {
 					// Attempt to reconnect to a running step
-					if reconnector, ok := r.engine.(backend.Reconnector); ok {
-						reconnectErr := reconnector.Reconnect(r.ctx, step, r.taskUUID)
-						if reconnectErr == nil {
-							logger.Info().Str("step", step.Name).Msg("reconnecting to existing step")
-							return r.execReconnected(step)
-						}
-						logger.Debug().Err(reconnectErr).Str("step", step.Name).Msg("cannot reconnect, re-executing step")
-					} else {
-						logger.Debug().Str("step", step.Name).Msg("backend does not support reconnection, re-executing step")
+					reconnectErr := r.engine.Reconnect(r.ctx, step, r.taskUUID)
+					if reconnectErr == nil {
+						logger.Info().Str("step", step.Name).Msg("reconnecting to existing step")
+						return r.execReconnected(step)
 					}
+					logger.Debug().Err(reconnectErr).Str("step", step.Name).Msg("cannot reconnect, re-executing step")
 				}
 
 				// Mark step as running in recovery state

pipeline/recovery.go

@@ -50,7 +50,7 @@ func NewRecoveryManager(client RecoveryClient, workflowID string, enabled bool)
 // On first run, creates recovery states for all steps.
 // On agent restart, loads existing states into cache.
 func (m *RecoveryManager) InitRecoveryState(ctx context.Context, config *backend.Config, timeoutSeconds int64) error {
-	if !m.enabled || m.client == nil {
+	if !m.enabled {
 		return nil
 	}
 
@@ -78,7 +78,7 @@ func (m *RecoveryManager) GetStepState(step *backend.Step) *rpc.RecoveryState {
 
 // MarkStepRunning marks a step as running.
 func (m *RecoveryManager) MarkStepRunning(ctx context.Context, step *backend.Step) error {
-	if !m.enabled || m.client == nil {
+	if !m.enabled {
 		return nil
 	}
 
@@ -87,7 +87,7 @@ func (m *RecoveryManager) MarkStepRunning(ctx context.Context, step *backend.Ste
 
 // MarkStepSuccess marks a step as successfully completed.
 func (m *RecoveryManager) MarkStepSuccess(ctx context.Context, step *backend.Step) error {
-	if !m.enabled || m.client == nil {
+	if !m.enabled {
 		return nil
 	}
 
@@ -96,7 +96,7 @@ func (m *RecoveryManager) MarkStepSuccess(ctx context.Context, step *backend.Ste
 
 // MarkStepFailed marks a step as failed.
 func (m *RecoveryManager) MarkStepFailed(ctx context.Context, step *backend.Step, exitCode int) error {
-	if !m.enabled || m.client == nil {
+	if !m.enabled {
 		return nil
 	}