Support for multiple trace context encodings

This is part merge-conflict-resolution, part refinement of #1775.
The main differences to #1775 are: specifying the encoding with a string flag
rather than bool, represented internally with a `Codec` interface, and some
tweaks to the tests.

Co-authored-by: Sam Park <[email protected]>

Author: DrJosh9000

agent/agent_configuration.go

@@ -61,5 +61,6 @@ type AgentConfiguration struct {
 	AcquireJob                 string
 	TracingBackend             string
 	TracingServiceName         string
+	TraceContextEncoding       string
 	DisableWarningsFor         []string
 }

agent/job_runner.go

@@ -502,6 +502,7 @@ func (r *JobRunner) createEnvironment(ctx context.Context) ([]string, error) {
 	env["BUILDKITE_STRICT_SINGLE_HOOKS"] = fmt.Sprintf("%t", r.conf.AgentConfiguration.StrictSingleHooks)
 	env["BUILDKITE_CANCEL_GRACE_PERIOD"] = strconv.Itoa(r.conf.AgentConfiguration.CancelGracePeriod)
 	env["BUILDKITE_SIGNAL_GRACE_PERIOD_SECONDS"] = strconv.Itoa(int(r.conf.AgentConfiguration.SignalGracePeriod / time.Second))
+	env["BUILDKITE_TRACE_CONTEXT_ENCODING"] = r.conf.AgentConfiguration.TraceContextEncoding
 
 	if r.conf.KubernetesExec {
 		env["BUILDKITE_KUBERNETES_EXEC"] = "true"

clicommand/agent_start.go

@@ -161,13 +161,14 @@ type AgentStartConfig struct {
 	TracingServiceName          string `cli:"tracing-service-name"`
 
 	// Global flags
-	Debug             bool     `cli:"debug"`
-	LogLevel          string   `cli:"log-level"`
-	NoColor           bool     `cli:"no-color"`
-	Experiments       []string `cli:"experiment" normalize:"list"`
-	Profile           string   `cli:"profile"`
-	StrictSingleHooks bool     `cli:"strict-single-hooks"`
-	KubernetesExec    bool     `cli:"kubernetes-exec"`
+	Debug                bool     `cli:"debug"`
+	LogLevel             string   `cli:"log-level"`
+	NoColor              bool     `cli:"no-color"`
+	Experiments          []string `cli:"experiment" normalize:"list"`
+	Profile              string   `cli:"profile"`
+	StrictSingleHooks    bool     `cli:"strict-single-hooks"`
+	KubernetesExec       bool     `cli:"kubernetes-exec"`
+	TraceContextEncoding string   `cli:"trace-context-encoding"`
 
 	// API config
 	DebugHTTP bool   `cli:"debug-http"`
@@ -689,6 +690,7 @@ var AgentStartCommand = cli.Command{
 		RedactedVars,
 		StrictSingleHooksFlag,
 		KubernetesExecFlag,
+		TraceContextEncodingFlag,
 
 		// Deprecated flags which will be removed in v4
 		cli.StringSliceFlag{
@@ -849,6 +851,10 @@ var AgentStartCommand = cli.Command{
 			return err
 		}
 
+		if _, err := tracetools.ParseEncoding(cfg.TraceContextEncoding); err != nil {
+			return fmt.Errorf("while parsing trace context encoding: %v", err)
+		}
+
 		mc := metrics.NewCollector(l, metrics.CollectorConfig{
 			Datadog:              cfg.MetricsDatadog,
 			DatadogHost:          cfg.MetricsDatadogHost,
@@ -946,6 +952,7 @@ var AgentStartCommand = cli.Command{
 			AcquireJob:                   cfg.AcquireJob,
 			TracingBackend:               cfg.TracingBackend,
 			TracingServiceName:           cfg.TracingServiceName,
+			TraceContextEncoding:         cfg.TraceContextEncoding,
 			VerificationFailureBehaviour: cfg.VerificationFailureBehavior,
 			KubernetesExec:               cfg.KubernetesExec,
 

clicommand/bootstrap.go

@@ -11,6 +11,7 @@ import (
 
 	"github.com/buildkite/agent/v3/internal/job"
 	"github.com/buildkite/agent/v3/process"
+	"github.com/buildkite/agent/v3/tracetools"
 	"github.com/urfave/cli"
 )
 
@@ -95,6 +96,7 @@ type BootstrapConfig struct {
 	RedactedVars                 []string `cli:"redacted-vars" normalize:"list"`
 	TracingBackend               string   `cli:"tracing-backend"`
 	TracingServiceName           string   `cli:"tracing-service-name"`
+	TraceContextEncoding         string   `cli:"trace-context-encoding"`
 	NoJobAPI                     bool     `cli:"no-job-api"`
 	DisableWarningsFor           []string `cli:"disable-warnings-for" normalize:"list"`
 	KubernetesExec               bool     `cli:"kubernetes-exec"`
@@ -382,6 +384,7 @@ var BootstrapCommand = cli.Command{
 		RedactedVars,
 		StrictSingleHooksFlag,
 		KubernetesExecFlag,
+		TraceContextEncodingFlag,
 	},
 	Action: func(c *cli.Context) error {
 		ctx := context.Background()
@@ -414,6 +417,11 @@ var BootstrapCommand = cli.Command{
 			return err
 		}
 
+		traceContextCodec, err := tracetools.ParseEncoding(cfg.TraceContextEncoding)
+		if err != nil {
+			return fmt.Errorf("while parsing trace context encoding: %v", err)
+		}
+
 		// Configure the bootstraper
 		bootstrap := job.New(job.ExecutorConfig{
 			AgentName:                    cfg.AgentName,
@@ -464,6 +472,7 @@ var BootstrapCommand = cli.Command{
 			Tag:                          cfg.Tag,
 			TracingBackend:               cfg.TracingBackend,
 			TracingServiceName:           cfg.TracingServiceName,
+			TraceContextCodec:            traceContextCodec,
 			JobAPI:                       !cfg.NoJobAPI,
 			DisabledWarnings:             cfg.DisableWarningsFor,
 			KubernetesExec:               cfg.KubernetesExec,

clicommand/global.go

@@ -117,6 +117,13 @@ var (
 			"*_CONNECTION_STRING",
 		},
 	}
+
+	TraceContextEncodingFlag = cli.StringFlag{
+		Name:   "trace-context-encoding",
+		Usage:  "Sets the inner encoding for BUILDKITE_TRACE_CONTEXT. Must be either json or gob",
+		Value:  "gob",
+		EnvVar: "BUILDKITE_TRACE_CONTEXT_ENCODING",
+	}
 )
 
 func globalFlags() []cli.Flag {

internal/job/config.go

@@ -8,6 +8,7 @@ import (
 
 	"github.com/buildkite/agent/v3/env"
 	"github.com/buildkite/agent/v3/process"
+	"github.com/buildkite/agent/v3/tracetools"
 )
 
 // Config provides the configuration for the job executor. Some of the keys are
@@ -165,6 +166,9 @@ type ExecutorConfig struct {
 	// Service name to use when reporting traces.
 	TracingServiceName string
 
+	// Encoding (within base64) for the trace context environment variable.
+	TraceContextCodec tracetools.Codec
+
 	// Whether to start the JobAPI
 	JobAPI bool
 

internal/job/executor.go

@@ -89,7 +89,7 @@ func (e *Executor) Run(ctx context.Context) (exitCode int) {
 		var err error
 		logger := shell.StderrLogger
 		logger.DisabledWarningIDs = e.DisabledWarnings
-		e.shell, err = shell.New(shell.WithLogger(logger))
+		e.shell, err = shell.New(shell.WithLogger(logger), shell.WithTraceContextCodec(e.TraceContextCodec))
 		if err != nil {
 			fmt.Printf("Error creating shell: %v", err)
 			return 1

internal/job/shell/shell.go

@@ -72,6 +72,9 @@ type Shell struct {
 
 	// Amount of time to wait between sending the InterruptSignal and SIGKILL
 	SignalGracePeriod time.Duration
+
+	// How to encode trace contexts.
+	traceContextCodec tracetools.Codec
 }
 
 type newShellOpt func(*Shell)
@@ -82,6 +85,12 @@ func WithLogger(l Logger) newShellOpt {
 	}
 }
 
+func WithTraceContextCodec(c tracetools.Codec) newShellOpt {
+	return func(s *Shell) {
+		s.traceContextCodec = c
+	}
+}
+
 // New returns a new Shell
 func New(opts ...newShellOpt) (*Shell, error) {
 	wd, err := os.Getwd()
@@ -90,10 +99,11 @@ func New(opts ...newShellOpt) (*Shell, error) {
 	}
 
 	shell := &Shell{
-		Logger: StderrLogger,
-		Env:    env.FromSlice(os.Environ()),
-		Writer: os.Stdout,
-		wd:     wd,
+		Logger:            StderrLogger,
+		Env:               env.FromSlice(os.Environ()),
+		Writer:            os.Stdout,
+		wd:                wd,
+		traceContextCodec: tracetools.CodecGob{},
 	}
 
 	for _, opt := range opts {
@@ -119,6 +129,7 @@ func (s *Shell) WithStdin(r io.Reader) *Shell {
 		wd:                s.wd,
 		InterruptSignal:   s.InterruptSignal,
 		SignalGracePeriod: s.SignalGracePeriod,
+		traceContextCodec: s.traceContextCodec,
 	}
 }
 
@@ -373,7 +384,7 @@ func (s *Shell) injectTraceCtx(ctx context.Context, env *env.Environment) {
 	if span == nil {
 		return
 	}
-	if err := tracetools.EncodeTraceContext(span, env.Dump()); err != nil {
+	if err := tracetools.EncodeTraceContext(span, env.Dump(), s.traceContextCodec); err != nil {
 		if s.Debug {
 			s.Logger.Warningf("Failed to encode trace context: %v", err)
 		}

internal/job/tracing.go

@@ -100,7 +100,7 @@ func (e *Executor) startTracingDatadog(ctx context.Context) (tracetools.Span, co
 // extractTraceCtx pulls encoded distributed tracing information from the env vars.
 // Note: This should match the injectTraceCtx code in shell.
 func (e *Executor) extractDDTraceCtx() opentracing.SpanContext {
-	sctx, err := tracetools.DecodeTraceContext(e.shell.Env.Dump())
+	sctx, err := tracetools.DecodeTraceContext(e.shell.Env.Dump(), e.ExecutorConfig.TraceContextCodec)
 	if err != nil {
 		// Return nil so a new span will be created
 		return nil

tracetools/propagate.go

@@ -4,6 +4,9 @@ import (
 	"bytes"
 	"encoding/base64"
 	"encoding/gob"
+	"encoding/json"
+	"fmt"
+	"io"
 
 	"github.com/opentracing/opentracing-go"
 	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
@@ -15,14 +18,14 @@ const EnvVarTraceContextKey = "BUILDKITE_TRACE_CONTEXT"
 
 // EncodeTraceContext will serialize and encode tracing data into a string and place
 // it into the given env vars map.
-func EncodeTraceContext(span opentracing.Span, env map[string]string) error {
+func EncodeTraceContext(span opentracing.Span, env map[string]string, codec Codec) error {
 	textmap := tracer.TextMapCarrier{}
 	if err := span.Tracer().Inject(span.Context(), opentracing.TextMap, &textmap); err != nil {
 		return err
 	}
 
-	buf := bytes.NewBuffer([]byte{})
-	enc := gob.NewEncoder(buf)
+	buf := bytes.NewBuffer(nil)
+	enc := codec.NewEncoder(buf)
 	if err := enc.Encode(textmap); err != nil {
 		return err
 	}
@@ -33,7 +36,7 @@ func EncodeTraceContext(span opentracing.Span, env map[string]string) error {
 
 // DecodeTraceContext will decode, deserialize, and extract the tracing data from the
 // given env var map.
-func DecodeTraceContext(env map[string]string) (opentracing.SpanContext, error) {
+func DecodeTraceContext(env map[string]string, codec Codec) (opentracing.SpanContext, error) {
 	s, has := env[EnvVarTraceContextKey]
 	if !has {
 		return nil, opentracing.ErrSpanContextNotFound
@@ -44,12 +47,49 @@ func DecodeTraceContext(env map[string]string) (opentracing.SpanContext, error)
 		return nil, err
 	}
 
-	buf := bytes.NewBuffer(contextBytes)
-	dec := gob.NewDecoder(buf)
+	dec := codec.NewDecoder(bytes.NewReader(contextBytes))
 	textmap := opentracing.TextMapCarrier{}
 	if err := dec.Decode(&textmap); err != nil {
 		return nil, err
 	}
 
 	return opentracing.GlobalTracer().Extract(opentracing.TextMap, textmap)
 }
+
+// Encoder impls can encode values. Decoder impls can decode values.
+type Encoder interface{ Encode(v any) error }
+type Decoder interface{ Decode(v any) error }
+
+// Codec implementations produce encoders/decoders.
+type Codec interface {
+	NewEncoder(io.Writer) Encoder
+	NewDecoder(io.Reader) Decoder
+	String() string
+}
+
+// CodecGob marshals and unmarshals with https://pkg.go.dev/encoding/gob.
+type CodecGob struct{}
+
+func (CodecGob) NewEncoder(w io.Writer) Encoder { return gob.NewEncoder(w) }
+func (CodecGob) NewDecoder(r io.Reader) Decoder { return gob.NewDecoder(r) }
+func (CodecGob) String() string                 { return "gob" }
+
+// CodecJSON marshals and unmarshals with https://pkg.go.dev/encoding/json.
+type CodecJSON struct{}
+
+func (CodecJSON) NewEncoder(w io.Writer) Encoder { return json.NewEncoder(w) }
+func (CodecJSON) NewDecoder(r io.Reader) Decoder { return json.NewDecoder(r) }
+func (CodecJSON) String() string                 { return "json" }
+
+// ParseEncoding converts an encoding to the associated codec.
+// An empty string is parsed as "gob".
+func ParseEncoding(encoding string) (Codec, error) {
+	switch encoding {
+	case "", "gob":
+		return CodecGob{}, nil
+	case "json":
+		return CodecJSON{}, nil
+	default:
+		return nil, fmt.Errorf("invalid encoding %q", encoding)
+	}
+}

tracetools/propagate_example_test.go

@@ -35,7 +35,7 @@ func ExampleEncodeTraceContext() {
 
 		// Now say we want to launch a child process.
 		// Prepare it's env vars. This will be the carrier for the tracing data.
-		if err := EncodeTraceContext(span, childEnv); err != nil {
+		if err := EncodeTraceContext(span, childEnv, CodecGob{}); err != nil {
 			fmt.Println("oops an error for parent process trace injection")
 		}
 		// Now childEnv will contain the encoded data set with the env var key.
@@ -58,7 +58,7 @@ func ExampleEncodeTraceContext() {
 		// Make sure tracing is setup the same way (same env var key)
 		// Normally you'd use os.Environ or similar here (the list of strings is
 		// supported). We're just reusing childEnv for test simplicity.
-		sctx, err := DecodeTraceContext(childEnv)
+		sctx, err := DecodeTraceContext(childEnv, CodecGob{})
 		if err != nil {
 			fmt.Println("oops an error for child process trace extraction")
 		} else {