batches: upload large exec artifacts

Batch executor jobs can produce stdout, stderr, and diffs large enough to bloat the final JSONL result stream. Upload oversized step artifacts from the existing `src batch exec` execution path instead of adding a separate command or user-facing auth flags.

`src batch exec` now discovers the executor job context from the environment, uses the existing executor job bearer token plus executor name auth model to stream large artifacts to the Batch Changes artifacts endpoint, and records returned references on `AfterStepResult`. Small outputs stay inline for compatibility, and per-step progress output is suppressed when artifact uploads are enabled so large logs are not duplicated into task-step JSONL events.

companian branch in sg: wb/proto-exec-upload

Test Plan:

- go test ./cmd/src ./internal/batches/executor ./internal/batches/ui github.com/sourcegraph/sourcegraph/lib/batches/execution

Author: burmudar

cmd/src/batch_exec.go

@@ -123,7 +123,11 @@ Examples:
 }
 
 func executeBatchSpecInWorkspaces(ctx context.Context, flags *executorModeFlags) (err error) {
-	ui := &ui.JSONLines{BinaryDiffs: flags.binaryDiffs}
+	artifactUploader, err := newBatchArtifactUploaderFromEnv(cfg.endpointURL)
+	if err != nil {
+		return err
+	}
+	ui := &ui.JSONLines{BinaryDiffs: flags.binaryDiffs, SuppressStepOutput: artifactUploader != nil}
 
 	// Ensure the temp dir exists.
 	tempDir := flags.tempDir
@@ -214,6 +218,9 @@ func executeBatchSpecInWorkspaces(ctx context.Context, flags *executorModeFlags)
 		UI:               taskExecUI.StepsExecutionUI(task),
 		ForceRoot:        !flags.runAsImageUser,
 		BinaryDiffs:      flags.binaryDiffs,
+
+		ArtifactUploader:             artifactUploader,
+		ArtifactUploadThresholdBytes: defaultArtifactUploadThresholdBytes,
 	}
 	results, err := executor.RunSteps(ctx, opts)
 

cmd/src/batch_exec_artifacts.go

@@ -0,0 +1,108 @@
+package main
+
+import (
+	"context"
+	"encoding/json"
+	"io"
+	"net/http"
+	"net/url"
+	"os"
+	"strconv"
+	"strings"
+
+	"github.com/sourcegraph/src-cli/internal/batches/executor"
+
+	"github.com/sourcegraph/sourcegraph/lib/batches/execution"
+	"github.com/sourcegraph/sourcegraph/lib/errors"
+)
+
+// defaultArtifactUploadThresholdBytes is 1 MiB.
+const defaultArtifactUploadThresholdBytes = 1 << 20
+
+type batchArtifactUploader struct {
+	endpointURL  *url.URL
+	jobID        int
+	jobToken     string
+	executorName string
+	client       *http.Client
+}
+
+var _ executor.ArtifactUploader = (*batchArtifactUploader)(nil)
+
+func newBatchArtifactUploaderFromEnv(endpointURL *url.URL) (*batchArtifactUploader, error) {
+	jobID, _ := strconv.Atoi(os.Getenv("SRC_EXECUTOR_JOB_ID"))
+	jobToken := os.Getenv("SRC_EXECUTOR_JOB_TOKEN")
+	executorName := os.Getenv("SRC_EXECUTOR_NAME")
+
+	configured := jobID != 0 || jobToken != "" || executorName != ""
+	if !configured {
+		return nil, nil
+	}
+	if jobID == 0 || jobToken == "" || executorName == "" {
+		return nil, errors.New("artifact upload requires job ID, job token, and executor name")
+	}
+
+	return &batchArtifactUploader{
+		endpointURL:  endpointURL,
+		jobID:        jobID,
+		jobToken:     jobToken,
+		executorName: executorName,
+		client:       http.DefaultClient,
+	}, nil
+}
+
+func (u *batchArtifactUploader) Upload(ctx context.Context, artifactKey string, r io.Reader) (execution.Artifact, error) {
+	if strings.Contains(artifactKey, "/") || strings.Contains(artifactKey, "\\") || strings.Contains(artifactKey, "..") {
+		return execution.Artifact{}, errors.Newf("invalid artifact key %q", artifactKey)
+	}
+	sizeReader := &artifactSizeReader{r: r}
+
+	url := u.endpointURL.JoinPath(
+		".executors",
+		"queue",
+		"batches",
+		"artifacts",
+		artifactKey,
+	)
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url.String(), sizeReader)
+	if err != nil {
+		return execution.Artifact{}, err
+	}
+	req.Header.Set("Authorization", "Bearer "+u.jobToken)
+	req.Header.Set("X-Sourcegraph-Executor-Name", u.executorName)
+	req.Header.Set("X-Sourcegraph-Job-ID", strconv.Itoa(u.jobID))
+
+	resp, err := u.client.Do(req)
+	if err != nil {
+		return execution.Artifact{}, err
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
+		body, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
+		return execution.Artifact{}, errors.Newf("artifact upload failed with status %s: %s", resp.Status, strings.TrimSpace(string(body)))
+	}
+
+	var ref execution.Artifact
+	if err := json.NewDecoder(resp.Body).Decode(&ref); err != nil {
+		return execution.Artifact{}, errors.Wrap(err, "decoding artifact upload response")
+	}
+	if ref.ObjectStorageKey == "" {
+		return execution.Artifact{}, errors.New("artifact upload response did not include an object storage key")
+	}
+	ref.Size = sizeReader.size
+	return ref, nil
+}
+
+type artifactSizeReader struct {
+	r    io.Reader
+	size int64
+}
+
+func (r *artifactSizeReader) Read(p []byte) (int, error) {
+	n, err := r.r.Read(p)
+	if n > 0 {
+		r.size += int64(n)
+	}
+	return n, err
+}

cmd/src/batch_exec_artifacts_test.go

@@ -0,0 +1,72 @@
+package main
+
+import (
+	"context"
+	"encoding/json"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"strings"
+	"testing"
+
+	"github.com/sourcegraph/sourcegraph/lib/batches/execution"
+)
+
+func TestBatchArtifactUploaderUploadAddsMetadata(t *testing.T) {
+	const artifactContents = "artifact contents"
+
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.Method != http.MethodPost {
+			t.Fatalf("unexpected method %q", r.Method)
+		}
+		if r.URL.Path != "/.executors/queue/batches/artifacts/stdout" {
+			t.Fatalf("unexpected path %q", r.URL.Path)
+		}
+		if got := r.Header.Get("Authorization"); got != "Bearer token" {
+			t.Fatalf("unexpected authorization header %q", got)
+		}
+		if got := r.Header.Get("X-Sourcegraph-Executor-Name"); got != "executor" {
+			t.Fatalf("unexpected executor header %q", got)
+		}
+		if got := r.Header.Get("X-Sourcegraph-Job-ID"); got != "42" {
+			t.Fatalf("unexpected job ID header %q", got)
+		}
+
+		body, err := io.ReadAll(r.Body)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if string(body) != artifactContents {
+			t.Fatalf("unexpected body %q", string(body))
+		}
+
+		w.Header().Set("Content-Type", "application/json")
+		_ = json.NewEncoder(w).Encode(execution.Artifact{ObjectStorageKey: "key"})
+	}))
+	t.Cleanup(server.Close)
+
+	endpointURL, err := url.Parse(server.URL)
+	if err != nil {
+		t.Fatal(err)
+	}
+	uploader := &batchArtifactUploader{
+		endpointURL:  endpointURL,
+		jobID:        42,
+		jobToken:     "token",
+		executorName: "executor",
+		client:       server.Client(),
+	}
+
+	ref, err := uploader.Upload(context.Background(), "stdout", strings.NewReader(artifactContents))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if ref.ObjectStorageKey != "key" {
+		t.Fatalf("unexpected object storage key %q", ref.ObjectStorageKey)
+	}
+	if ref.Size != int64(len(artifactContents)) {
+		t.Fatalf("unexpected size %d", ref.Size)
+	}
+}

internal/batches/executor/artifacts.go

@@ -0,0 +1,144 @@
+package executor
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"io"
+	"math"
+	"os"
+
+	"github.com/sourcegraph/sourcegraph/lib/batches/execution"
+	"github.com/sourcegraph/sourcegraph/lib/errors"
+)
+
+const maxInlineArtifactSize = math.MaxInt64
+
+type ArtifactUploader interface {
+	Upload(ctx context.Context, artifactKey string, r io.Reader) (execution.Artifact, error)
+}
+
+type stepOutput struct {
+	stdout *artifactOutput
+	stderr *artifactOutput
+}
+
+func newStepOutput(dir string, threshold int64) (*stepOutput, error) {
+	stdout, err := newArtifactOutput(dir, "stdout-*", threshold)
+	if err != nil {
+		return nil, err
+	}
+	stderr, err := newArtifactOutput(dir, "stderr-*", threshold)
+	if err != nil {
+		stdout.cleanup()
+		return nil, err
+	}
+	return &stepOutput{stdout: stdout, stderr: stderr}, nil
+}
+
+func (o *stepOutput) cleanup() {
+	o.stdout.cleanup()
+	o.stderr.cleanup()
+}
+
+type artifactOutput struct {
+	file      *os.File
+	buf       bytes.Buffer
+	size      int64
+	threshold int64
+}
+
+func newArtifactOutput(dir, pattern string, threshold int64) (*artifactOutput, error) {
+	file, err := os.CreateTemp(dir, pattern)
+	if err != nil {
+		return nil, errors.Wrap(err, "creating artifact output file")
+	}
+	return &artifactOutput{file: file, threshold: threshold}, nil
+}
+
+func (o *artifactOutput) writer() io.Writer { return o }
+
+func (o *artifactOutput) Write(p []byte) (int, error) {
+	n, err := o.file.Write(p)
+	o.size += int64(n)
+	if o.size <= o.threshold {
+		_, _ = o.buf.Write(p[:n])
+	}
+	return n, err
+}
+
+func (o *artifactOutput) inline() string {
+	if o.size > o.threshold {
+		return ""
+	}
+	return o.buf.String()
+}
+
+func (o *artifactOutput) shouldUpload(threshold int64) bool {
+	return o.size > threshold
+}
+
+func (o *artifactOutput) reader() (io.Reader, error) {
+	if _, err := o.file.Seek(0, io.SeekStart); err != nil {
+		return nil, errors.Wrap(err, "rewinding artifact output file")
+	}
+	return o.file, nil
+}
+
+func (o *artifactOutput) cleanup() {
+	if o.file == nil {
+		return
+	}
+	name := o.file.Name()
+	_ = o.file.Close()
+	_ = os.Remove(name)
+	o.file = nil
+}
+
+func uploadStepArtifacts(ctx context.Context, uploader ArtifactUploader, threshold int64, stepIndex int, result *execution.AfterStepResult, output *stepOutput) error {
+	defer output.cleanup()
+	result.Artifacts = make(map[string]execution.Artifact)
+
+	if output.stdout.shouldUpload(threshold) {
+		ref, err := uploadArtifactOutput(ctx, uploader, artifactKey(stepIndex, execution.ArtifactStdout), output.stdout)
+		if err != nil {
+			return err
+		}
+		result.Artifacts[execution.ArtifactStdout] = ref
+	}
+
+	if output.stderr.shouldUpload(threshold) {
+		ref, err := uploadArtifactOutput(ctx, uploader, artifactKey(stepIndex, execution.ArtifactStderr), output.stderr)
+		if err != nil {
+			return err
+		}
+		result.Artifacts[execution.ArtifactStderr] = ref
+	}
+
+	if int64(len(result.Diff)) > threshold {
+		ref, err := uploader.Upload(ctx, artifactKey(stepIndex, execution.ArtifactDiff), bytes.NewReader(result.Diff))
+		if err != nil {
+			return errors.Wrap(err, "uploading diff artifact")
+		}
+		result.Artifacts[execution.ArtifactDiff] = ref
+		result.Diff = nil
+	}
+
+	return nil
+}
+
+func uploadArtifactOutput(ctx context.Context, uploader ArtifactUploader, key string, output *artifactOutput) (execution.Artifact, error) {
+	reader, err := output.reader()
+	if err != nil {
+		return execution.Artifact{}, err
+	}
+	ref, err := uploader.Upload(ctx, key, reader)
+	if err != nil {
+		return execution.Artifact{}, errors.Wrapf(err, "uploading %s artifact", key)
+	}
+	return ref, nil
+}
+
+func artifactKey(stepIndex int, name string) string {
+	return fmt.Sprintf("step-%d-%s", stepIndex, name)
+}