Merge pull request #2316 from buildkite/experiments-in-context

Store experiments in contexts

agent/agent_worker.go

@@ -651,7 +651,7 @@ func (a *AgentWorker) RunJob(ctx context.Context, acceptResponse *api.Job) error
 	})
 
 	// Now that we've got a job to do, we can start it.
-	jr, err := NewJobRunner(a.logger, a.apiClient, JobRunnerConfig{
+	jr, err := NewJobRunner(ctx, a.logger, a.apiClient, JobRunnerConfig{
 		Job:                acceptResponse,
 		JWKS:               a.agentConfiguration.JobVerificationJWKS,
 		Debug:              a.debug,

agent/artifact_uploader.go

@@ -15,7 +15,7 @@ import (
 	"time"
 
 	"github.com/buildkite/agent/v3/api"
-	"github.com/buildkite/agent/v3/experiments"
+	"github.com/buildkite/agent/v3/internal/experiments"
 	"github.com/buildkite/agent/v3/internal/mime"
 	"github.com/buildkite/agent/v3/logger"
 	"github.com/buildkite/agent/v3/pool"
@@ -73,7 +73,7 @@ func NewArtifactUploader(l logger.Logger, ac APIClient, c ArtifactUploaderConfig
 
 func (a *ArtifactUploader) Upload(ctx context.Context) error {
 	// Create artifact structs for all the files we need to upload
-	artifacts, err := a.Collect()
+	artifacts, err := a.Collect(ctx)
 	if err != nil {
 		return fmt.Errorf("collecting artifacts: %w", err)
 	}
@@ -107,7 +107,7 @@ func isDir(path string) bool {
 	return fi.IsDir()
 }
 
-func (a *ArtifactUploader) Collect() (artifacts []*api.Artifact, err error) {
+func (a *ArtifactUploader) Collect(ctx context.Context) (artifacts []*api.Artifact, err error) {
 	wd, err := os.Getwd()
 	if err != nil {
 		return nil, fmt.Errorf("getting working directory: %w", err)
@@ -182,7 +182,7 @@ func (a *ArtifactUploader) Collect() (artifacts []*api.Artifact, err error) {
 				return nil, fmt.Errorf("resolving relative path for file %s: %w", file, err)
 			}
 
-			if experiments.IsEnabled(experiments.NormalisedUploadPaths) {
+			if experiments.IsEnabled(ctx, experiments.NormalisedUploadPaths) {
 				// Convert any Windows paths to Unix/URI form
 				path = filepath.ToSlash(path)
 			}

agent/artifact_uploader_test.go

@@ -1,14 +1,15 @@
 package agent
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"path/filepath"
 	"strings"
 	"testing"
 
 	"github.com/buildkite/agent/v3/api"
-	"github.com/buildkite/agent/v3/experiments"
+	"github.com/buildkite/agent/v3/internal/experiments"
 	"github.com/buildkite/agent/v3/logger"
 	"github.com/stretchr/testify/assert"
 )
@@ -24,7 +25,7 @@ func findArtifact(artifacts []*api.Artifact, search string) *api.Artifact {
 }
 
 func TestCollect(t *testing.T) {
-	// t.Parallel() cannot be used with experiments.Enable
+	ctx := context.Background()
 
 	wd, _ := os.Getwd()
 	root := filepath.Join(wd, "..")
@@ -100,24 +101,16 @@ func TestCollect(t *testing.T) {
 	// path.Join function instead (which uses Unix/URI-style path separators,
 	// regardless of platform)
 
-	experimentKey := experiments.NormalisedUploadPaths
-	experimentPrev := experiments.IsEnabled(experimentKey)
-	defer func() {
-		if experimentPrev {
-			experiments.Enable(experimentKey)
-		} else {
-			experiments.Disable(experimentKey)
-		}
-	}()
-	experiments.Disable(experimentKey)
-	artifactsWithoutExperimentEnabled, err := uploader.Collect()
+	ctxExpEnabled, _ := experiments.Enable(ctx, experiments.NormalisedUploadPaths)
+	ctxExpDisabled := experiments.Disable(ctx, experiments.NormalisedUploadPaths)
+
+	artifactsWithoutExperimentEnabled, err := uploader.Collect(ctxExpDisabled)
 	if err != nil {
 		t.Fatalf("[normalised-upload-paths disabled] uploader.Collect() error = %v", err)
 	}
 	assert.Equal(t, 5, len(artifactsWithoutExperimentEnabled))
 
-	experiments.Enable(experimentKey)
-	artifactsWithExperimentEnabled, err := uploader.Collect()
+	artifactsWithExperimentEnabled, err := uploader.Collect(ctxExpEnabled)
 	if err != nil {
 		t.Fatalf("[normalised-upload-paths enabled] uploader.Collect() error = %v", err)
 	}
@@ -167,6 +160,8 @@ func TestCollect(t *testing.T) {
 }
 
 func TestCollectThatDoesntMatchAnyFiles(t *testing.T) {
+	ctx := context.Background()
+
 	wd, _ := os.Getwd()
 	root := filepath.Join(wd, "..")
 	os.Chdir(root)
@@ -181,7 +176,7 @@ func TestCollectThatDoesntMatchAnyFiles(t *testing.T) {
 		}, ";"),
 	})
 
-	artifacts, err := uploader.Collect()
+	artifacts, err := uploader.Collect(ctx)
 	if err != nil {
 		t.Fatalf("uploader.Collect() error = %v", err)
 	}
@@ -190,6 +185,8 @@ func TestCollectThatDoesntMatchAnyFiles(t *testing.T) {
 }
 
 func TestCollectWithSomeGlobsThatDontMatchAnything(t *testing.T) {
+	ctx := context.Background()
+
 	wd, _ := os.Getwd()
 	root := filepath.Join(wd, "..")
 	os.Chdir(root)
@@ -203,7 +200,7 @@ func TestCollectWithSomeGlobsThatDontMatchAnything(t *testing.T) {
 		}, ";"),
 	})
 
-	artifacts, err := uploader.Collect()
+	artifacts, err := uploader.Collect(ctx)
 	if err != nil {
 		t.Fatalf("uploader.Collect() error = %v", err)
 	}
@@ -214,6 +211,8 @@ func TestCollectWithSomeGlobsThatDontMatchAnything(t *testing.T) {
 }
 
 func TestCollectWithSomeGlobsThatDontMatchAnythingFollowingSymlinks(t *testing.T) {
+	ctx := context.Background()
+
 	wd, _ := os.Getwd()
 	root := filepath.Join(wd, "..")
 	os.Chdir(root)
@@ -229,7 +228,7 @@ func TestCollectWithSomeGlobsThatDontMatchAnythingFollowingSymlinks(t *testing.T
 		GlobResolveFollowSymlinks: true,
 	})
 
-	artifacts, err := uploader.Collect()
+	artifacts, err := uploader.Collect(ctx)
 	if err != nil {
 		t.Fatalf("uploader.Collect() error = %v", err)
 	}
@@ -240,6 +239,8 @@ func TestCollectWithSomeGlobsThatDontMatchAnythingFollowingSymlinks(t *testing.T
 }
 
 func TestCollectWithDuplicateMatches(t *testing.T) {
+	ctx := context.Background()
+
 	wd, _ := os.Getwd()
 	root := filepath.Join(wd, "..")
 	os.Chdir(root)
@@ -252,7 +253,7 @@ func TestCollectWithDuplicateMatches(t *testing.T) {
 		}, ";"),
 	})
 
-	artifacts, err := uploader.Collect()
+	artifacts, err := uploader.Collect(ctx)
 	if err != nil {
 		t.Fatalf("uploader.Collect() error = %v", err)
 	}
@@ -274,6 +275,8 @@ func TestCollectWithDuplicateMatches(t *testing.T) {
 }
 
 func TestCollectWithDuplicateMatchesFollowingSymlinks(t *testing.T) {
+	ctx := context.Background()
+
 	wd, _ := os.Getwd()
 	root := filepath.Join(wd, "..")
 	os.Chdir(root)
@@ -287,7 +290,7 @@ func TestCollectWithDuplicateMatchesFollowingSymlinks(t *testing.T) {
 		GlobResolveFollowSymlinks: true,
 	})
 
-	artifacts, err := uploader.Collect()
+	artifacts, err := uploader.Collect(ctx)
 	if err != nil {
 		t.Fatalf("uploader.Collect() error = %v", err)
 	}
@@ -310,6 +313,8 @@ func TestCollectWithDuplicateMatchesFollowingSymlinks(t *testing.T) {
 }
 
 func TestCollectMatchesUploadSymlinks(t *testing.T) {
+	ctx := context.Background()
+
 	wd, _ := os.Getwd()
 	root := filepath.Join(wd, "..")
 	os.Chdir(root)
@@ -322,7 +327,7 @@ func TestCollectMatchesUploadSymlinks(t *testing.T) {
 		UploadSkipSymlinks: true,
 	})
 
-	artifacts, err := uploader.Collect()
+	artifacts, err := uploader.Collect(ctx)
 	if err != nil {
 		t.Fatalf("uploader.Collect() error = %v", err)
 	}

agent/integration/job_runner_integration_test.go

@@ -1,6 +1,7 @@
 package integration
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"strconv"
@@ -13,6 +14,7 @@ import (
 
 func TestPreBootstrapHookRefusesJob(t *testing.T) {
 	t.Parallel()
+	ctx := context.Background()
 
 	hooksDir, err := os.MkdirTemp("", "bootstrap-hooks")
 	if err != nil {
@@ -45,7 +47,7 @@ func TestPreBootstrapHookRefusesJob(t *testing.T) {
 	mb.Expect().NotCalled() // The bootstrap won't be called, as the pre-bootstrap hook failed
 	defer mb.CheckAndClose(t)
 
-	runJob(t, testRunJobConfig{
+	runJob(t, ctx, testRunJobConfig{
 		job:           j,
 		server:        server,
 		agentCfg:      agent.AgentConfiguration{HooksPath: hooksDir},
@@ -65,6 +67,7 @@ func TestPreBootstrapHookRefusesJob(t *testing.T) {
 
 func TestJobRunner_WhenBootstrapExits_ItSendsTheExitStatusToTheAPI(t *testing.T) {
 	t.Parallel()
+	ctx := context.Background()
 
 	exits := []int{0, 1, 2, 3}
 	for _, exit := range exits {
@@ -89,7 +92,7 @@ func TestJobRunner_WhenBootstrapExits_ItSendsTheExitStatusToTheAPI(t *testing.T)
 			server := e.server("my-job-id")
 			defer server.Close()
 
-			runJob(t, testRunJobConfig{
+			runJob(t, ctx, testRunJobConfig{
 				job:           j,
 				server:        server,
 				agentCfg:      agent.AgentConfiguration{},
@@ -106,6 +109,7 @@ func TestJobRunner_WhenBootstrapExits_ItSendsTheExitStatusToTheAPI(t *testing.T)
 
 func TestJobRunner_WhenJobHasToken_ItOverridesAccessToken(t *testing.T) {
 	t.Parallel()
+	ctx := context.Background()
 
 	jobToken := "actually-llamas-are-only-okay"
 
@@ -133,7 +137,7 @@ func TestJobRunner_WhenJobHasToken_ItOverridesAccessToken(t *testing.T) {
 	server := e.server("my-job-id")
 	defer server.Close()
 
-	runJob(t, testRunJobConfig{
+	runJob(t, ctx, testRunJobConfig{
 		job:           j,
 		server:        server,
 		agentCfg:      agent.AgentConfiguration{},
@@ -145,6 +149,7 @@ func TestJobRunner_WhenJobHasToken_ItOverridesAccessToken(t *testing.T) {
 // Maybe that the job runner pulls the access token from the API client? but that's all handled in the `runJob` helper...
 func TestJobRunnerPassesAccessTokenToBootstrap(t *testing.T) {
 	t.Parallel()
+	ctx := context.Background()
 
 	j := &api.Job{
 		ID:                 "my-job-id",
@@ -169,7 +174,7 @@ func TestJobRunnerPassesAccessTokenToBootstrap(t *testing.T) {
 	server := e.server("my-job-id")
 	defer server.Close()
 
-	runJob(t, testRunJobConfig{
+	runJob(t, ctx, testRunJobConfig{
 		job:           j,
 		server:        server,
 		agentCfg:      agent.AgentConfiguration{},
@@ -179,6 +184,7 @@ func TestJobRunnerPassesAccessTokenToBootstrap(t *testing.T) {
 
 func TestJobRunnerIgnoresPipelineChangesToProtectedVars(t *testing.T) {
 	t.Parallel()
+	ctx := context.Background()
 
 	j := &api.Job{
 		ID:                 "my-job-id",
@@ -204,7 +210,7 @@ func TestJobRunnerIgnoresPipelineChangesToProtectedVars(t *testing.T) {
 	server := e.server("my-job-id")
 	defer server.Close()
 
-	runJob(t, testRunJobConfig{
+	runJob(t, ctx, testRunJobConfig{
 		job:           j,
 		server:        server,
 		agentCfg:      agent.AgentConfiguration{CommandEval: true},

agent/integration/job_verification_integration_test.go

@@ -1,6 +1,7 @@
 package integration
 
 import (
+	"context"
 	"fmt"
 	"strings"
 	"testing"
@@ -91,6 +92,7 @@ var (
 
 func TestJobVerification(t *testing.T) {
 	t.Parallel()
+	ctx := context.Background()
 
 	cases := []struct {
 		name                     string
@@ -227,7 +229,7 @@ func TestJobVerification(t *testing.T) {
 			defer mb.CheckAndClose(t)
 
 			tc.job.Step = signStep(t, pipelineUploadEnv, tc.job.Step, tc.signingKey)
-			runJob(t, testRunJobConfig{
+			runJob(t, ctx, testRunJobConfig{
 				job:              &tc.job,
 				server:           server,
 				agentCfg:         tc.agentConf,

agent/integration/test_helpers.go

@@ -47,7 +47,7 @@ type testRunJobConfig struct {
 	verificationJWKS jwk.Set
 }
 
-func runJob(t *testing.T, cfg testRunJobConfig) {
+func runJob(t *testing.T, ctx context.Context, cfg testRunJobConfig) {
 	t.Helper()
 
 	l := logger.Discard
@@ -64,7 +64,7 @@ func runJob(t *testing.T, cfg testRunJobConfig) {
 		Token:    "llamasrock",
 	})
 
-	jr, err := agent.NewJobRunner(l, client, agent.JobRunnerConfig{
+	jr, err := agent.NewJobRunner(ctx, l, client, agent.JobRunnerConfig{
 		Job:                cfg.job,
 		JWKS:               cfg.verificationJWKS,
 		AgentConfiguration: cfg.agentCfg,

agent/job_runner.go

@@ -12,7 +12,7 @@ import (
 	"time"
 
 	"github.com/buildkite/agent/v3/api"
-	"github.com/buildkite/agent/v3/experiments"
+	"github.com/buildkite/agent/v3/internal/experiments"
 	"github.com/buildkite/agent/v3/internal/job/shell"
 	"github.com/buildkite/agent/v3/kubernetes"
 	"github.com/buildkite/agent/v3/logger"
@@ -162,7 +162,7 @@ type jobAPI interface {
 var _ jobRunner = (*JobRunner)(nil)
 
 // Initializes the job runner
-func NewJobRunner(l logger.Logger, apiClient APIClient, conf JobRunnerConfig) (jobRunner, error) {
+func NewJobRunner(ctx context.Context, l logger.Logger, apiClient APIClient, conf JobRunnerConfig) (jobRunner, error) {
 	r := &JobRunner{
 		logger:    l,
 		conf:      conf,
@@ -220,7 +220,7 @@ func NewJobRunner(l logger.Logger, apiClient APIClient, conf JobRunnerConfig) (j
 		r.envFile = file
 	}
 
-	env, err := r.createEnvironment()
+	env, err := r.createEnvironment(ctx)
 	if err != nil {
 		return nil, err
 	}
@@ -333,7 +333,7 @@ func NewJobRunner(l logger.Logger, apiClient APIClient, conf JobRunnerConfig) (j
 	processEnv := append(os.Environ(), env...)
 
 	// The process that will run the bootstrap script
-	if experiments.IsEnabled(experiments.KubernetesExec) {
+	if experiments.IsEnabled(ctx, experiments.KubernetesExec) {
 		containerCount, err := strconv.Atoi(os.Getenv("BUILDKITE_CONTAINER_COUNT"))
 		if err != nil {
 			return nil, fmt.Errorf("failed to parse BUILDKITE_CONTAINER_COUNT: %w", err)
@@ -392,7 +392,7 @@ func (r *JobRunner) normalizeVerificationBehavior(behavior string) (string, erro
 }
 
 // Creates the environment variables that will be used in the process and writes a flat environment file
-func (r *JobRunner) createEnvironment() ([]string, error) {
+func (r *JobRunner) createEnvironment(ctx context.Context) ([]string, error) {
 	// Create a clone of our jobs environment. We'll then set the
 	// environment variables provided by the agent, which will override any
 	// sent by Buildkite. The variables below should always take
@@ -476,7 +476,7 @@ func (r *JobRunner) createEnvironment() ([]string, error) {
 	env["BUILDKITE_GIT_CLEAN_FLAGS"] = r.conf.AgentConfiguration.GitCleanFlags
 	env["BUILDKITE_GIT_MIRRORS_LOCK_TIMEOUT"] = fmt.Sprintf("%d", r.conf.AgentConfiguration.GitMirrorsLockTimeout)
 	env["BUILDKITE_SHELL"] = r.conf.AgentConfiguration.Shell
-	env["BUILDKITE_AGENT_EXPERIMENT"] = strings.Join(experiments.Enabled(), ",")
+	env["BUILDKITE_AGENT_EXPERIMENT"] = strings.Join(experiments.Enabled(ctx), ",")
 	env["BUILDKITE_REDACTED_VARS"] = strings.Join(r.conf.AgentConfiguration.RedactedVars, ",")
 	env["BUILDKITE_STRICT_SINGLE_HOOKS"] = fmt.Sprintf("%t", r.conf.AgentConfiguration.StrictSingleHooks)