PB-5: improve exit code propogation for non-command containers

Author: zhming0

internal/controller/scheduler/fail_job.go

@@ -19,6 +19,13 @@ import (
 	"k8s.io/client-go/kubernetes"
 )
 
+// FailureInfo contains data about a job failure for reporting to Buildkite.
+type FailureInfo struct {
+	Message  string
+	ExitCode int32
+	Reason   string
+}
+
 // acquireAndFailForObject figures out how to fail the BK job corresponding to
 // the k8s object (a pod or job) by inspecting the object's labels.
 func acquireAndFailForObject(
@@ -27,7 +34,7 @@ func acquireAndFailForObject(
 	k8sClient kubernetes.Interface,
 	cfg *config.Config,
 	obj metav1.Object,
-	message string,
+	failureCtx FailureInfo,
 ) error {
 	agentToken, err := fetchAgentToken(ctx, logger, k8sClient, obj.GetNamespace(), cfg.AgentTokenSecret)
 	if err != nil {
@@ -45,7 +52,7 @@ func acquireAndFailForObject(
 	tags := agenttags.TagsFromLabels(labels)
 	opts := cfg.AgentConfig.ControllerOptions()
 
-	if err := acquireAndFail(ctx, logger, agentToken, cfg.JobPrefix, jobUUID, tags, message, opts...); err != nil {
+	if err := acquireAndFail(ctx, logger, agentToken, cfg.JobPrefix, jobUUID, tags, failureCtx, opts...); err != nil {
 		logger.Error("failed to acquire and fail the job on Buildkite", zap.Error(err))
 		return err
 	}
@@ -61,7 +68,7 @@ func acquireAndFail(
 	jobPrefix string,
 	jobUUID string,
 	tags []string,
-	message string,
+	failureCtx FailureInfo,
 	options ...agentcore.ControllerOption,
 ) error {
 	opts := append([]agentcore.ControllerOption{
@@ -89,13 +96,24 @@ func acquireAndFail(
 		return fmt.Errorf("starting job: %w", err)
 	}
 
-	if err := jctr.WriteLog(ctx, message); err != nil {
+	if err := jctr.WriteLog(ctx, failureCtx.Message); err != nil {
 		zapLogger.Error("writing log", zap.Error(err))
 		return fmt.Errorf("writing log: %w", err)
 	}
 
 	var ignoreAgentInDispatches *bool
-	if err := jctr.Finish(ctx, agentcore.ProcessExit{Status: 1, SignalReason: agent.SignalReasonStackError}, ignoreAgentInDispatches); err != nil {
+
+	status := failureCtx.ExitCode
+	if status == 0 {
+		status = 1
+	}
+	reason := failureCtx.Reason
+	if reason == "" {
+		// By default, we consider all failure triggered by the stack controller are stack errors
+		reason = agent.SignalReasonStackError
+	}
+
+	if err := jctr.Finish(ctx, agentcore.ProcessExit{Status: int(status), SignalReason: reason}, ignoreAgentInDispatches); err != nil {
 		zapLogger.Error("finishing job", zap.Error(err))
 		return fmt.Errorf("finishing job: %w", err)
 	}

internal/controller/scheduler/job_watcher.go

@@ -8,6 +8,7 @@ import (
 
 	"github.com/buildkite/agent-stack-k8s/v2/internal/controller/config"
 	"github.com/buildkite/agent-stack-k8s/v2/internal/controller/model"
+	"github.com/buildkite/agent/v3/agent"
 
 	"github.com/google/uuid"
 	"github.com/jedib0t/go-pretty/v6/table"
@@ -231,7 +232,12 @@ func (w *jobWatcher) fetchEvents(ctx context.Context, log *zap.Logger, kjob *bat
 }
 
 func (w *jobWatcher) failJob(ctx context.Context, log *zap.Logger, kjob *batchv1.Job, message string) {
-	if err := acquireAndFailForObject(ctx, log, w.k8s, w.cfg, kjob, message); err != nil {
+	failureCtx := FailureInfo{
+		Message: message,
+		// We can know almost all failures triggered by job watcher are stack related error.
+		Reason: agent.SignalReasonStackError,
+	}
+	if err := acquireAndFailForObject(ctx, log, w.k8s, w.cfg, kjob, failureCtx); err != nil {
 		// Maybe the job was cancelled in the meantime?
 		log.Error("Could not fail Buildkite job", zap.Error(err))
 		jobWatcherBuildkiteJobFailErrorsCounter.Inc()

internal/controller/scheduler/pod_watcher.go

@@ -11,6 +11,7 @@ import (
 
 	"github.com/buildkite/agent-stack-k8s/v2/api"
 	"github.com/buildkite/agent-stack-k8s/v2/internal/controller/config"
+	"github.com/buildkite/agent/v3/agent"
 	"github.com/buildkite/roko"
 
 	"github.com/google/uuid"
@@ -297,6 +298,7 @@ func (w *podWatcher) failOnInitContainerFailure(ctx context.Context, log *zap.Lo
 	log.Debug("Checking pod for failed init containers")
 
 	containerFails := make(map[string]*corev1.ContainerStateTerminated)
+	var lastFailExitCode int32
 
 	// If any init container fails, whether it's one we added specifically to
 	// check for pull failure or not, the pod won't run.
@@ -306,6 +308,7 @@ func (w *podWatcher) failOnInitContainerFailure(ctx context.Context, log *zap.Lo
 			continue
 		}
 		containerFails[containerStatus.Name] = term
+		lastFailExitCode = term.ExitCode
 	}
 
 	if len(containerFails) == 0 {
@@ -322,7 +325,12 @@ func (w *podWatcher) failOnInitContainerFailure(ctx context.Context, log *zap.Lo
 	// probably shouldn't interfere.
 	log.Info("One or more init containers failed. Failing.")
 	message := w.formatInitContainerFails(containerFails)
-	if err := acquireAndFailForObject(ctx, log, w.k8s, w.cfg, pod, message); err != nil {
+	failureCtx := FailureInfo{
+		Message:  message,
+		ExitCode: lastFailExitCode,
+		Reason:   agent.SignalReasonStackError,
+	}
+	if err := acquireAndFailForObject(ctx, log, w.k8s, w.cfg, pod, failureCtx); err != nil {
 		// Maybe the job was cancelled in the meantime?
 		log.Error("Could not fail Buildkite job", zap.Error(err))
 		podWatcherBuildkiteJobFailErrorsCounter.Inc()
@@ -477,7 +485,11 @@ func (w *podWatcher) failForImageFailure(ctx context.Context, log *zap.Logger, f
 		// We can acquire it and fail it ourselves.
 		log.Info("One or more job containers are waiting too long for images. Failing.")
 		message := w.formatImagePullFailureMessage(statuses)
-		if err := acquireAndFailForObject(ctx, log, w.k8s, w.cfg, pod, message); err != nil {
+		failureCtx := FailureInfo{
+			Message: message,
+			// Do we have a better status code to report here?
+		}
+		if err := acquireAndFailForObject(ctx, log, w.k8s, w.cfg, pod, failureCtx); err != nil {
 			podWatcherBuildkiteJobFailErrorsCounter.Inc()
 			// Maybe the job was acquired by an agent in the meantime?
 			// Maybe the job was cancelled in the meantime?

internal/controller/scheduler/scheduler.go

@@ -21,6 +21,7 @@ import (
 	"github.com/buildkite/agent-stack-k8s/v2/internal/version"
 	"github.com/buildkite/roko"
 
+	"github.com/buildkite/agent/v3/agent"
 	"github.com/buildkite/agent/v3/clicommand"
 
 	"github.com/distribution/reference"
@@ -1272,7 +1273,14 @@ func (w *worker) failJob(ctx context.Context, inputs buildInputs, message string
 	}
 
 	opts := w.cfg.AgentConfig.ControllerOptions()
-	if err := acquireAndFail(ctx, w.logger, agentToken, w.cfg.JobPrefix, inputs.uuid, inputs.agentQueryRules, message, opts...); err != nil {
+	failureCtx := FailureInfo{
+		Message: message,
+		// In scheduler worker, often these failure are technically users error in YAML.
+		// Using agent refuse seems to be reasonable as that represent that agent won't execute these yaml.
+		// Deserve more discussion though..
+		Reason: agent.SignalReasonAgentRefused,
+	}
+	if err := acquireAndFail(ctx, w.logger, agentToken, w.cfg.JobPrefix, inputs.uuid, inputs.agentQueryRules, failureCtx, opts...); err != nil {
 		w.logger.Error("failed to acquire and fail the job on Buildkite", zap.Error(err))
 		schedulerBuildkiteJobFailErrorsCounter.Inc()
 		return err