Merge branch 'ajwalker/refactor-steps-connector' into 'main'

Refactor Connector to allow setup before connection

See merge request https://gitlab.com/gitlab-org/gitlab-runner/-/merge_requests/6359

Merged-by: Arran Walker <ajwalker@gitlab.com>
Approved-by: Cameron Swords <cswords@gitlab.com>
Reviewed-by: Cameron Swords <cswords@gitlab.com>
Reviewed-by: GitLab Duo <gitlab-duo@gitlab.com>
Reviewed-by: Emmanuel 326 <nyariboemmanuel8@gmail.com>

Author: saracen

commands/steps/steps.go

@@ -27,6 +27,8 @@ import (
 
 const (
 	SubCommandName = "steps"
+
+	readyMessage = "step-runner is ready."
 )
 
 type IOStreams struct {
@@ -69,6 +71,7 @@ func Bootstrap(destination string) error {
 	return os.Chmod(destination, 0o755)
 }
 
+//nolint:gocognit
 func Serve(ctx context.Context, sockPath string, ioStreams IOStreams, cmdAndArgs ...string) error {
 	listener, err := net.ListenUnix("unix", api.SocketAddr(sockPath))
 	if err != nil {
@@ -87,6 +90,9 @@ func Serve(ctx context.Context, sockPath string, ioStreams IOStreams, cmdAndArgs
 	srv := grpc.NewServer()
 	proto.RegisterStepRunnerServer(srv, service)
 
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
 	wg, ctx := errgroup.WithContext(ctx)
 
 	go func() {
@@ -102,8 +108,13 @@ func Serve(ctx context.Context, sockPath string, ioStreams IOStreams, cmdAndArgs
 		return nil
 	})
 
+	fmt.Fprintln(os.Stderr, readyMessage)
+
 	if len(cmdAndArgs) > 0 {
 		wg.Go(func() error {
+			// on script exit, we cancel() so that the step-runner serve also terminates
+			defer cancel()
+
 			stdin := bufio.NewReader(ioStreams.Stdin)
 
 			stdinCheck := make(chan error, 1)

common/build.go

@@ -418,6 +418,7 @@ func (b *Build) StartBuild(
 	return nil
 }
 
+//nolint:gocognit
 func (b *Build) executeStepStage(ctx context.Context, connector steps.Connector, buildStage BuildStage, req []schema.Step) error {
 	if ctx.Err() != nil {
 		return ctx.Err()
@@ -453,30 +454,46 @@ func (b *Build) executeStepStage(ctx context.Context, connector steps.Connector,
 				Variables:  b.GetAllVariables(),
 			}
 
-			err := steps.Execute(ctx, connector, info, req, stdout)
-			if err != nil {
-				berr := &BuildError{Inner: err}
-
-				var cserr *steps.ClientStatusError
-				if errors.As(err, &cserr) {
-					switch cserr.Status.State {
-					case client.StateUnspecified:
-						berr.FailureReason = UnknownFailure
-					case client.StateFailure:
-						berr.FailureReason = ScriptFailure
-					}
-				}
-
-				return berr
-			}
-
-			return err
+			return wrapStepStageErr(steps.Execute(ctx, connector, info, req, stdout))
 		},
 	}
 
 	return section.Execute(&b.logger)
 }
 
+func wrapStepStageErr(err error) error {
+	if err == nil {
+		return nil
+	}
+
+	if errors.Is(err, steps.ErrNoStepRunnerButOkay) {
+		return nil
+	}
+
+	berr := &BuildError{Inner: err}
+
+	var cserr *steps.ClientStatusError
+	if errors.As(err, &cserr) {
+		switch cserr.Status.State {
+		case client.StateUnspecified:
+			berr.FailureReason = UnknownFailure
+		case client.StateFailure:
+			berr.FailureReason = ScriptFailure
+		}
+	}
+
+	// hack: for now, we parse the exit code from the error response
+	// later we might want to introduce a proper exit code from the step-runner
+	// https://gitlab.com/gitlab-org/step-runner/-/work_items/349
+	if _, code, ok := strings.Cut(err.Error(), "exit status"); ok {
+		if exitCode, err := strconv.Atoi(strings.TrimSpace(code)); err == nil {
+			berr.ExitCode = exitCode
+		}
+	}
+
+	return berr
+}
+
 //nolint:gocognit
 func (b *Build) executeStage(ctx context.Context, buildStage BuildStage, executor Executor) error {
 	if b.UseNativeSteps() {

common/mocks.go

@@ -883,19 +883,20 @@ func (e *executor) createContainerConfig(
 		OpenStdin:    true,
 		StdinOnce:    true,
 		Entrypoint:   e.overwriteEntrypoint(&imageDefinition),
+		Env:          e.Build.GetAllVariables().StringList(),
 	}
 
-	// We should never really need to set the environment variables on the container,
-	// as these are exported via the abstract shell.
-	//
-	// Adding these variables interferes with steps. Given this situation, when native
-	// steps is enabled, we no longer add the env vars to the container.
-	if !e.Build.UseNativeSteps() {
-		config.Env = e.Build.GetAllVariables().StringList()
-	}
-
-	// user config should only be set in build containers
+	//nolint:nestif
 	if containerType == buildContainerType {
+		if e.Build.UseNativeSteps() {
+			config.Cmd = append([]string{bootstrappedBinary, "steps", "serve"}, config.Cmd...)
+
+			// Environment variables interferes with steps. Given this situation, when
+			// native steps are enabled, we no longer add the env vars to the container.
+			config.Env = nil
+		}
+
+		// user config should only be set in build containers
 		if user, err := e.getBuildContainerUser(imageDefinition); err != nil {
 			return nil, err
 		} else {

executors/docker/docker.go

@@ -204,18 +204,12 @@ func (s *commandExecutor) requestBuildContainer() (*container.InspectResponse, e
 		return s.buildContainer, nil
 	}
 
-	// Overwrite the build container command if using steps
-	cmd := s.BuildShell.DockerCommand
-	if s.Build.UseNativeSteps() {
-		cmd = append([]string{bootstrappedBinary, "steps", "serve"}, s.BuildShell.DockerCommand...)
-	}
-
 	var err error
 	s.buildContainer, err = s.createContainer(
 		buildContainerType,
 		s.Build.Image,
 		[]string{},
-		newDefaultContainerConfigurator(&s.executor, buildContainerType, s.Build.Image, cmd, []string{}),
+		newDefaultContainerConfigurator(&s.executor, buildContainerType, s.Build.Image, s.BuildShell.DockerCommand, []string{}),
 	)
 	if err != nil {
 		return nil, err

executors/docker/docker_command.go

@@ -171,7 +171,7 @@ func (e *machineExecutor) GetMetricsSelector() string {
 	return refereed.GetMetricsSelector()
 }
 
-func (s *machineExecutor) Connect(ctx context.Context) (io.ReadWriteCloser, error) {
+func (s *machineExecutor) Connect(ctx context.Context) (func() (io.ReadWriteCloser, error), error) {
 	if connector, ok := s.executor.(steps.Connector); ok {
 		return connector.Connect(ctx)
 	}

executors/docker/machine/machine.go

@@ -12,6 +12,7 @@ import (
 	"gitlab.com/gitlab-org/gitlab-runner/common"
 	"gitlab.com/gitlab-org/gitlab-runner/helpers/docker"
 	"gitlab.com/gitlab-org/gitlab-runner/session/terminal"
+	"gitlab.com/gitlab-org/gitlab-runner/steps"
 )
 
 func getRunnerConfig() *common.RunnerConfig {
@@ -135,7 +136,7 @@ func TestMachineCredentialsUsage(t *testing.T) {
 type mockDockerExecutor struct {
 	*common.MockExecutor
 	*terminal.MockInteractiveTerminal
-	*common.MockConnector
+	*steps.MockConnector
 }
 
 func TestMachineExecutor_WithoutInteractiveTerminal(t *testing.T) {
@@ -177,7 +178,7 @@ func TestMachineExecutor_WithInteractiveTerminal(t *testing.T) {
 func TestMachineExecutor_Connect(t *testing.T) {
 	mock := mockDockerExecutor{
 		MockExecutor:  common.NewMockExecutor(t),
-		MockConnector: common.NewMockConnector(t),
+		MockConnector: steps.NewMockConnector(t),
 	}
 	e := machineExecutor{
 		executor: &mock,