Fix pipeline cancellation status handling and step state synchronization

agent/logger.go

@@ -36,6 +36,7 @@ func (r *Runner) createLogger(_logger zerolog.Logger, uploads *sync.WaitGroup, w
 			Logger()
 
 		uploads.Add(1)
+		defer uploads.Done()
 
 		var secrets []string
 		for _, secret := range workflow.Config.Secrets {
@@ -50,8 +51,6 @@ func (r *Runner) createLogger(_logger zerolog.Logger, uploads *sync.WaitGroup, w
 		}
 
 		logger.Debug().Msg("log stream copied, close ...")
-		uploads.Done()
-
 		return nil
 	}
 }

agent/rpc/client_grpc.go

@@ -148,26 +148,27 @@ func (c *client) Next(ctx context.Context, filter rpc.Filter) (*rpc.Workflow, er
 	return w, nil
 }
 
-// Wait blocks until the workflow is complete.
-func (c *client) Wait(ctx context.Context, workflowID string) (err error) {
+// Wait blocks until the workflow with the given ID is marked as completed or canceled by the server.
+func (c *client) Wait(ctx context.Context, workflowID string) (canceled bool, err error) {
 	retry := c.newBackOff()
 	req := new(proto.WaitRequest)
 	req.Id = workflowID
 	for {
-		_, err = c.client.Wait(ctx, req)
+		resp, err := c.client.Wait(ctx, req)
 		if err == nil {
-			break
+			// wait block was released normally as expected by server
+			return resp.GetCanceled(), nil
 		}
 
 		switch status.Code(err) {
 		case codes.Canceled:
 			if ctx.Err() != nil {
 				// expected as context was canceled
 				log.Debug().Err(err).Msgf("grpc error: wait(): context canceled")
-				return nil
+				return false, nil
 			}
 			log.Error().Err(err).Msgf("grpc error: wait(): code: %v", status.Code(err))
-			return err
+			return false, err
 		case
 			codes.Aborted,
 			codes.DataLoss,
@@ -178,16 +179,15 @@ func (c *client) Wait(ctx context.Context, workflowID string) (err error) {
 			log.Warn().Err(err).Msgf("grpc error: wait(): code: %v", status.Code(err))
 		default:
 			log.Error().Err(err).Msgf("grpc error: wait(): code: %v", status.Code(err))
-			return err
+			return false, err
 		}
 
 		select {
 		case <-time.After(retry.NextBackOff()):
 		case <-ctx.Done():
-			return ctx.Err()
+			return false, ctx.Err()
 		}
 	}
-	return nil
 }
 
 // Init signals the workflow is initialized.
@@ -199,6 +199,7 @@ func (c *client) Init(ctx context.Context, workflowID string, state rpc.Workflow
 	req.State.Started = state.Started
 	req.State.Finished = state.Finished
 	req.State.Error = state.Error
+	req.State.Canceled = state.Canceled
 	for {
 		_, err = c.client.Init(ctx, req)
 		if err == nil {
@@ -238,7 +239,7 @@ func (c *client) Init(ctx context.Context, workflowID string, state rpc.Workflow
 	return nil
 }
 
-// Done signals the workflow is complete.
+// Done let agent signal to server the workflow has stopped.
 func (c *client) Done(ctx context.Context, workflowID string, state rpc.WorkflowState) (err error) {
 	retry := c.newBackOff()
 	req := new(proto.DoneRequest)
@@ -247,6 +248,7 @@ func (c *client) Done(ctx context.Context, workflowID string, state rpc.Workflow
 	req.State.Started = state.Started
 	req.State.Finished = state.Finished
 	req.State.Error = state.Error
+	req.State.Canceled = state.Canceled
 	for {
 		_, err = c.client.Done(ctx, req)
 		if err == nil {
@@ -330,7 +332,7 @@ func (c *client) Extend(ctx context.Context, workflowID string) (err error) {
 	return nil
 }
 
-// Update updates the workflow state.
+// Update let agent updates the step state at the server.
 func (c *client) Update(ctx context.Context, workflowID string, state rpc.StepState) (err error) {
 	retry := c.newBackOff()
 	req := new(proto.UpdateRequest)
@@ -342,6 +344,7 @@ func (c *client) Update(ctx context.Context, workflowID string, state rpc.StepSt
 	req.State.Exited = state.Exited
 	req.State.ExitCode = int32(state.ExitCode)
 	req.State.Error = state.Error
+	req.State.Canceled = state.Canceled
 	for {
 		_, err = c.client.Update(ctx, req)
 		if err == nil {

agent/runner.go

@@ -20,7 +20,6 @@ import (
 	"errors"
 	"fmt"
 	"sync"
-	"sync/atomic"
 	"time"
 
 	"github.com/rs/zerolog/log"
@@ -51,6 +50,7 @@ func NewRunner(workEngine rpc.Peer, f rpc.Filter, h string, state *State, backen
 	}
 }
 
+// Run executes a workflow using a backend, tracks its state and reports the state back to the server.
 func (r *Runner) Run(runnerCtx, shutdownCtx context.Context) error {
 	log.Debug().Msg("request next execution")
 
@@ -90,34 +90,32 @@ func (r *Runner) Run(runnerCtx, shutdownCtx context.Context) error {
 
 	// Workflow execution context.
 	// This context is the SINGLE source of truth for cancellation.
-	workflowCtx, cancel := context.WithTimeout(ctxMeta, timeout)
-	defer cancel()
+	workflowCtx, _ := context.WithTimeout(ctxMeta, timeout) //nolint:govet
+	workflowCtx, cancelWorkflowCtx := context.WithCancelCause(workflowCtx)
+	defer cancelWorkflowCtx(nil)
 
-	// Handle SIGTERM (k8s, docker, system shutdown)
+	// Add sigterm support for internal context.
+	// Required to be able to terminate the running workflow by external signals.
 	workflowCtx = utils.WithContextSigtermCallback(workflowCtx, func() {
 		logger.Error().Msg("received sigterm termination signal")
-		cancel()
+		// WithContextSigtermCallback would cancel the context too, but  we want our own custom error
+		cancelWorkflowCtx(pipeline.ErrCancel)
 	})
 
-	// canceled indicates whether the workflow was canceled remotely (UI/API).
-	// Must be atomic because it is written from a goroutine and read later.
-	var canceled atomic.Bool
-
 	// Listen for remote cancel events (UI / API).
 	// When canceled, we MUST cancel the workflow context
-	// so that pipeline execution and backend processes stop immediately.
+	// so that workflow execution stop immediately.
 	go func() {
-		logger.Debug().Msg("listening for cancel signal")
-
-		if err := r.client.Wait(workflowCtx, workflow.ID); err != nil {
-			logger.Warn().Err(err).Msg("cancel signal received from server")
-
-			// Mark workflow as canceled (thread-safe)
-			canceled.Store(true)
+		logger.Debug().Msg("start listening for server side cancel signal")
 
-			// Propagate cancellation to pipeline + backend
-			cancel()
+		if canceled, err := r.client.Wait(workflowCtx, workflow.ID); err != nil {
+			logger.Error().Err(err).Msg("server returned unexpected err while waiting for workflow to finish run")
+			cancelWorkflowCtx(err)
 		} else {
+			if canceled {
+				logger.Debug().Err(err).Msg("server side cancel signal received")
+				cancelWorkflowCtx(pipeline.ErrCancel)
+			}
 			// Wait returned without error, meaning the workflow finished normally
 			logger.Debug().Msg("cancel listener exited normally")
 		}
@@ -143,9 +141,13 @@ 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("workflow initialization failed")
-		// TODO: should we return here?
+		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
 	}
 
 	var uploads sync.WaitGroup
@@ -167,19 +169,18 @@ func (r *Runner) Run(runnerCtx, shutdownCtx context.Context) error {
 
 	state.Finished = time.Now().Unix()
 
-	// Normalize cancellation error
-	if errors.Is(err, pipeline.ErrCancel) || canceled.Load() {
-		canceled.Store(true)
-		err = pipeline.ErrCancel
-	}
-
 	if err != nil {
 		state.Error = err.Error()
+		if errors.Is(err, pipeline.ErrCancel) {
+			state.Canceled = true
+			// cleanup joined error messages
+			state.Error = pipeline.ErrCancel.Error()
+		}
 	}
 
 	logger.Debug().
 		Str("error", state.Error).
-		Bool("canceled", canceled.Load()).
+		Bool("canceled", state.Canceled).
 		Msg("workflow finished")
 
 	// Ensure all logs/traces are uploaded before finishing
@@ -195,6 +196,8 @@ func (r *Runner) Run(runnerCtx, shutdownCtx context.Context) error {
 
 	if err := r.client.Done(doneCtx, workflow.ID, state); err != nil {
 		logger.Error().Err(err).Msg("failed to update workflow status")
+	} else {
+		logger.Debug().Msg("signaling workflow stopped done")
 	}
 
 	return nil

agent/tracer.go

@@ -16,6 +16,7 @@ package agent
 
 import (
 	"context"
+	"errors"
 	"runtime"
 	"strconv"
 	"sync"
@@ -30,6 +31,7 @@ import (
 func (r *Runner) createTracer(ctxMeta context.Context, uploads *sync.WaitGroup, logger zerolog.Logger, workflow *rpc.Workflow) pipeline.TraceFunc {
 	return func(state *pipeline.State) error {
 		uploads.Add(1)
+		defer uploads.Done()
 
 		stepLogger := logger.With().
 			Str("image", state.Pipeline.Step.Image).
@@ -43,12 +45,15 @@ func (r *Runner) createTracer(ctxMeta context.Context, uploads *sync.WaitGroup,
 			StepUUID: state.Pipeline.Step.UUID,
 			Exited:   state.Process.Exited,
 			ExitCode: state.Process.ExitCode,
-			Started:  time.Now().Unix(), // TODO: do not do this
-			Finished: time.Now().Unix(),
+			Started:  state.Process.Started,
+			Canceled: errors.Is(state.Process.Error, pipeline.ErrCancel),
 		}
 		if state.Process.Error != nil {
 			stepState.Error = state.Process.Error.Error()
 		}
+		if state.Process.Exited {
+			stepState.Finished = time.Now().Unix()
+		}
 
 		defer func() {
 			stepLogger.Debug().Msg("update step status")
@@ -60,7 +65,6 @@ func (r *Runner) createTracer(ctxMeta context.Context, uploads *sync.WaitGroup,
 			}
 
 			stepLogger.Debug().Msg("update step status complete")
-			uploads.Done()
 		}()
 		if state.Process.Exited {
 			return nil

pipeline/backend/docker/docker.go

@@ -250,29 +250,11 @@ func (e *docker) StartStep(ctx context.Context, step *backend.Step, taskUUID str
 	return e.client.ContainerStart(ctx, containerName, container.StartOptions{})
 }
 
-// WaitStep waits for a step container to exit.
-//
-// When the context is canceled, the container is immediately killed to prevent
-// orphaned containers from continuing to run after agent shutdown.
 func (e *docker) WaitStep(ctx context.Context, step *backend.Step, taskUUID string) (*backend.State, error) {
-	log := log.Logger.With().
-		Str("taskUUID", taskUUID).
-		Str("stepUUID", step.UUID).
-		Logger()
-
+	log := log.Logger.With().Str("taskUUID", taskUUID).Str("stepUUID", step.UUID).Logger()
 	log.Trace().Msgf("wait for step %s", step.Name)
 
 	containerName := toContainerName(step)
-	done := make(chan struct{})
-
-	// Ensure container is killed if context is canceled (SIGTERM / pipeline cancel)
-	go func() {
-		select {
-		case <-ctx.Done():
-			_ = e.client.ContainerKill(context.Background(), containerName, "9") //nolint:contextcheck
-		case <-done:
-		}
-	}()
 
 	wait, errC := e.client.ContainerWait(ctx, containerName, "")
 	select {
@@ -282,9 +264,6 @@ func (e *docker) WaitStep(ctx context.Context, step *backend.Step, taskUUID stri
 		log.Trace().Msgf("ContainerWait returned with err: %v", err)
 	}
 
-	// Stop cancellation watcher
-	close(done)
-
 	info, err := e.client.ContainerInspect(ctx, containerName)
 	if err != nil {
 		return nil, err

pipeline/backend/local/local.go

@@ -171,9 +171,18 @@ func (e *local) StartStep(ctx context.Context, step *types.Step, taskUUID string
 	}
 }
 
-func (e *local) WaitStep(_ context.Context, step *types.Step, taskUUID string) (*types.State, error) {
+func (e *local) WaitStep(ctx context.Context, step *types.Step, taskUUID string) (*types.State, error) {
 	log.Trace().Str("taskUUID", taskUUID).Msgf("wait for step %s", step.Name)
 
+	stepState := &types.State{
+		Exited: true,
+	}
+
+	if err := ctx.Err(); err != nil {
+		stepState.Error = err
+		return stepState, nil
+	}
+
 	state, err := e.getStepState(taskUUID, step.UUID)
 	if err != nil {
 		return nil, err
@@ -183,10 +192,6 @@ func (e *local) WaitStep(_ context.Context, step *types.Step, taskUUID string) (
 		return nil, errors.New("exec: step command not set up")
 	}
 
-	stepState := &types.State{
-		Exited: true,
-	}
-
 	// normally we use cmd.Wait() to wait for *exec.Cmd, but cmd.StdoutPipe() tells us not
 	// as Wait() would close the io pipe even if not all logs where read and send back
 	// so we have to do use the underlying functions

pipeline/backend/types/state.go

@@ -16,11 +16,14 @@ package types
 
 // State defines a container state.
 type State struct {
+	// Unix start time
+	Started int64 `json:"started"`
 	// Container exit code
 	ExitCode int `json:"exit_code"`
 	// Container exited, true or false
 	Exited bool `json:"exited"`
 	// Container is oom killed, true or false
+	// TODO (6024): well known errors as string enum into ./errors.go
 	OOMKilled bool `json:"oom_killed"`
 	// Container error
 	Error error