Add CPU core detection.

Author: LawnGnome

cmd/src/batch_common.go

@@ -23,6 +23,7 @@ import (
 
 	"github.com/sourcegraph/src-cli/internal/api"
 	"github.com/sourcegraph/src-cli/internal/batches"
+	"github.com/sourcegraph/src-cli/internal/batches/docker"
 	"github.com/sourcegraph/src-cli/internal/batches/executor"
 	"github.com/sourcegraph/src-cli/internal/batches/graphql"
 	"github.com/sourcegraph/src-cli/internal/batches/service"
@@ -121,8 +122,8 @@ func newBatchExecuteFlags(flagSet *flag.FlagSet, workspaceExecution bool, cacheD
 	)
 
 	flagSet.IntVar(
-		&caf.parallelism, "j", runtime.GOMAXPROCS(0),
-		"The maximum number of parallel jobs. Default is GOMAXPROCS.",
+		&caf.parallelism, "j", 0,
+		"The maximum number of parallel jobs. Default (or 0) is the number of CPU cores available to Docker, or GOMAXPROCS if Docker cannot report its number of cores.",
 	)
 	flagSet.DurationVar(
 		&caf.timeout, "timeout", 60*time.Minute,
@@ -275,7 +276,11 @@ func executeBatchSpec(ctx context.Context, ui ui.ExecUI, opts executeBatchSpecOp
 		return err
 	}
 
-	if err := checkExecutable("docker", "version"); err != nil {
+	// In the past, we checked `docker version`, but now we retrieve the number
+	// of CPUs, since we need that anyway and it performs the same check (is
+	// Docker working _at all_?).
+	parallelism, err := getBatchParallelism(ctx, opts.flags.parallelism)
+	if err != nil {
 		return err
 	}
 
@@ -307,7 +312,7 @@ func executeBatchSpec(ctx context.Context, ui ui.ExecUI, opts executeBatchSpecOp
 	if len(batchSpec.Steps) > 0 {
 		ui.PreparingContainerImages()
 		images, err := svc.EnsureDockerImages(
-			ctx, batchSpec.Steps, opts.flags.parallelism,
+			ctx, batchSpec.Steps, parallelism,
 			ui.PreparingContainerImagesProgress,
 		)
 		if err != nil {
@@ -354,7 +359,7 @@ func executeBatchSpec(ctx context.Context, ui ui.ExecUI, opts executeBatchSpecOp
 		Cache:           executor.NewDiskCache(opts.flags.cacheDir),
 		SkipErrors:      opts.flags.skipErrors,
 		CleanArchives:   opts.flags.cleanArchives,
-		Parallelism:     opts.flags.parallelism,
+		Parallelism:     parallelism,
 		Timeout:         opts.flags.timeout,
 		KeepLogs:        opts.flags.keepLogs,
 		TempDir:         opts.flags.tempDir,
@@ -386,7 +391,7 @@ func executeBatchSpec(ctx context.Context, ui ui.ExecUI, opts executeBatchSpecOp
 	}
 	ui.CheckingCacheSuccess(len(specs), len(uncachedTasks))
 
-	taskExecUI := ui.ExecutingTasks(*verbose, opts.flags.parallelism)
+	taskExecUI := ui.ExecutingTasks(*verbose, parallelism)
 	freshSpecs, logFiles, execErr := coord.ExecuteAndBuildSpecs(ctx, batchSpec, uncachedTasks, taskExecUI)
 	// Add external changeset specs.
 	importedSpecs, importErr := svc.CreateImportChangesetSpecs(ctx, batchSpec)
@@ -533,3 +538,25 @@ func contextCancelOnInterrupt(parent context.Context) (context.Context, func())
 		ctxCancel()
 	}
 }
+
+func getBatchParallelism(ctx context.Context, flag int) (int, error) {
+	if flag > 0 {
+		return flag, nil
+	}
+
+	ncpu, err := docker.NCPU(ctx)
+	var terr docker.TimeoutError
+	if errors.As(err, &terr) {
+		return 0, err
+	} else if err != nil {
+		// In the case of errors from Docker itself, we want to fall back to
+		// GOMAXPROCS, since it's possible Docker just doesn't have access to
+		// the CPU core count (either due to permissions, or being too old).
+		//
+		// It would obviously be better if we had a global logger available to
+		// log this.
+		return runtime.GOMAXPROCS(0), nil
+	}
+
+	return ncpu, nil
+}

cmd/src/batch_exec.go

@@ -92,7 +92,12 @@ func executeBatchSpecInWorkspaces(ctx context.Context, ui *ui.JSONLines, opts ex
 	if err := checkExecutable("git", "version"); err != nil {
 		return err
 	}
-	if err := checkExecutable("docker", "version"); err != nil {
+
+	// In the past, we checked `docker version`, but now we retrieve the number
+	// of CPUs, since we need that anyway and it performs the same check (is
+	// Docker working _at all_?).
+	parallelism, err := getBatchParallelism(ctx, opts.flags.parallelism)
+	if err != nil {
 		return err
 	}
 
@@ -112,7 +117,7 @@ func executeBatchSpecInWorkspaces(ctx context.Context, ui *ui.JSONLines, opts ex
 	if len(input.Steps) > 0 {
 		ui.PreparingContainerImages()
 		images, err := svc.EnsureDockerImages(
-			ctx, input.Steps, opts.flags.parallelism,
+			ctx, input.Steps, parallelism,
 			ui.PreparingContainerImagesProgress,
 		)
 		if err != nil {
@@ -138,7 +143,7 @@ func executeBatchSpecInWorkspaces(ctx context.Context, ui *ui.JSONLines, opts ex
 		Cache:         &executor.ServerSideCache{Writer: ui},
 		SkipErrors:    opts.flags.skipErrors,
 		CleanArchives: opts.flags.cleanArchives,
-		Parallelism:   opts.flags.parallelism,
+		Parallelism:   parallelism,
 		Timeout:       opts.flags.timeout,
 		KeepLogs:      opts.flags.keepLogs,
 		TempDir:       opts.flags.tempDir,
@@ -154,7 +159,7 @@ func executeBatchSpecInWorkspaces(ctx context.Context, ui *ui.JSONLines, opts ex
 		return err
 	}
 
-	taskExecUI := ui.ExecutingTasks(*verbose, opts.flags.parallelism)
+	taskExecUI := ui.ExecutingTasks(*verbose, parallelism)
 	err = coord.Execute(ctx, tasks, taskExecUI)
 	if err == nil || opts.flags.skipErrors {
 		if err == nil {

internal/batches/docker/info.go

@@ -0,0 +1,38 @@
+package docker
+
+import (
+	"bytes"
+	"context"
+	"runtime"
+	"strconv"
+
+	"github.com/sourcegraph/sourcegraph/lib/errors"
+
+	"github.com/sourcegraph/src-cli/internal/exec"
+)
+
+// NCPU returns the number of CPU cores available to Docker.
+func NCPU(ctx context.Context) (int, error) {
+	ncpu := runtime.GOMAXPROCS(0)
+
+	dctx, cancel, err := withFastCommandContext(ctx)
+	if err != nil {
+		return ncpu, err
+	}
+	defer cancel()
+
+	args := []string{"info", "--format", "{{ .NCPU }}"}
+	out, err := exec.CommandContext(dctx, "docker", args...).CombinedOutput()
+	if errors.Is(errors.Cause(err), context.DeadlineExceeded) {
+		return ncpu, newFastCommandTimeoutError(dctx, args...)
+	} else if err != nil {
+		return ncpu, err
+	}
+
+	dcpu, err := strconv.Atoi(string(bytes.TrimSpace(out)))
+	if err != nil {
+		return ncpu, err
+	}
+
+	return dcpu, nil
+}