From 69a8caa3bec0b95d5b3608dce6cfd6428c11e0cd Mon Sep 17 00:00:00 2001
From: Tonis Tiigi <tonistiigi@gmail.com>
Date: Tue, 8 Jun 2021 20:01:18 -0700
Subject: [PATCH] add current tracing context detection and exec propagation

Signed-off-by: Tonis Tiigi <tonistiigi@gmail.com>
(cherry picked from commit bc9a83144c83e9fd78007b7bfe92e8082c59d40e)
---
 executor/oci/spec.go           |  3 ++
 util/appcontext/appcontext.go  |  7 +++-
 util/appcontext/register.go    | 14 +++++++
 util/tracing/env/traceenv.go   | 54 +++++++++++++++++++++++++++
 util/tracing/exec/traceexec.go | 68 ++++++++++++++++++++++++++++++++++
 5 files changed, 145 insertions(+), 1 deletion(-)
 create mode 100644 util/appcontext/register.go
 create mode 100644 util/tracing/env/traceenv.go
 create mode 100644 util/tracing/exec/traceexec.go

diff --git a/executor/oci/spec.go b/executor/oci/spec.go
index 800031081..a01ae6c83 100644
--- a/executor/oci/spec.go
+++ b/executor/oci/spec.go
@@ -15,6 +15,7 @@ import (
 	"github.com/moby/buildkit/executor"
 	"github.com/moby/buildkit/snapshot"
 	"github.com/moby/buildkit/util/network"
+	traceexec "github.com/moby/buildkit/util/tracing/exec"
 	specs "github.com/opencontainers/runtime-spec/specs-go"
 	"github.com/opencontainers/selinux/go-selinux"
 	"github.com/pkg/errors"
@@ -76,6 +77,8 @@ func GenerateSpec(ctx context.Context, meta executor.Meta, mounts []executor.Mou
 		hostname = meta.Hostname
 	}
 
+	meta.Env = append(meta.Env, traceexec.Environ(ctx)...)
+
 	opts = append(opts,
 		oci.WithProcessArgs(meta.Args...),
 		oci.WithEnv(meta.Env),
diff --git a/util/appcontext/appcontext.go b/util/appcontext/appcontext.go
index e74da2cdb..3c67f9236 100644
--- a/util/appcontext/appcontext.go
+++ b/util/appcontext/appcontext.go
@@ -22,7 +22,12 @@ func Context() context.Context {
 		const exitLimit = 3
 		retries := 0
 
-		ctx, cancel := context.WithCancel(context.Background())
+		ctx := context.Background()
+		for _, f := range inits {
+			ctx = f(ctx)
+		}
+
+		ctx, cancel := context.WithCancel(ctx)
 		appContextCache = ctx
 
 		go func() {
diff --git a/util/appcontext/register.go b/util/appcontext/register.go
new file mode 100644
index 000000000..22f9e8365
--- /dev/null
+++ b/util/appcontext/register.go
@@ -0,0 +1,14 @@
+package appcontext
+
+import (
+	"context"
+)
+
+type Initializer func(context.Context) context.Context
+
+var inits []Initializer
+
+// Register stores a new context initializer that runs on app context creation
+func Register(f Initializer) {
+	inits = append(inits, f)
+}
diff --git a/util/tracing/env/traceenv.go b/util/tracing/env/traceenv.go
new file mode 100644
index 000000000..1c4e2f448
--- /dev/null
+++ b/util/tracing/env/traceenv.go
@@ -0,0 +1,54 @@
+package detect
+
+import (
+	"context"
+	"os"
+
+	"github.com/moby/buildkit/util/appcontext"
+	"go.opentelemetry.io/otel/propagation"
+)
+
+const (
+	traceparentHeader = "traceparent"
+	tracestateHeader  = "tracestate"
+)
+
+func init() {
+	appcontext.Register(initContext)
+}
+
+func initContext(ctx context.Context) context.Context {
+	// previously defined in https://github.com/open-telemetry/opentelemetry-swift/blob/4ea467ed4b881d7329bf2254ca7ed7f2d9d6e1eb/Sources/OpenTelemetrySdk/Trace/Propagation/EnvironmentContextPropagator.swift#L14-L15
+	parent := os.Getenv("OTEL_TRACE_PARENT")
+	state := os.Getenv("OTEL_TRACE_STATE")
+
+	if parent == "" {
+		return ctx
+	}
+
+	tc := propagation.TraceContext{}
+	return tc.Extract(ctx, &textMap{parent: parent, state: state})
+}
+
+type textMap struct {
+	parent string
+	state  string
+}
+
+func (tm *textMap) Get(key string) string {
+	switch key {
+	case traceparentHeader:
+		return tm.parent
+	case tracestateHeader:
+		return tm.state
+	default:
+		return ""
+	}
+}
+
+func (tm *textMap) Set(key string, value string) {
+}
+
+func (tm *textMap) Keys() []string {
+	return []string{traceparentHeader, tracestateHeader}
+}
diff --git a/util/tracing/exec/traceexec.go b/util/tracing/exec/traceexec.go
new file mode 100644
index 000000000..80020fbac
--- /dev/null
+++ b/util/tracing/exec/traceexec.go
@@ -0,0 +1,68 @@
+package detect
+
+import (
+	"context"
+
+	"go.opentelemetry.io/otel/propagation"
+)
+
+const (
+	traceparentHeader = "traceparent"
+	tracestateHeader  = "tracestate"
+)
+
+// Environ returns list of environment variables that need to be sent to the child process
+// in order for it to pick up cross-process tracing from same state.
+func Environ(ctx context.Context) []string {
+	var tm textMap
+	tc := propagation.TraceContext{}
+	tc.Inject(ctx, &tm)
+
+	var env []string
+
+	// previously defined in https://github.com/open-telemetry/opentelemetry-swift/blob/4ea467ed4b881d7329bf2254ca7ed7f2d9d6e1eb/Sources/OpenTelemetrySdk/Trace/Propagation/EnvironmentContextPropagator.swift#L14-L15
+	if tm.parent != "" {
+		env = append(env, "OTEL_TRACE_PARENT="+tm.parent)
+	}
+	if tm.state != "" {
+		env = append(env, "OTEL_TRACE_STATE="+tm.state)
+	}
+
+	return env
+}
+
+type textMap struct {
+	parent string
+	state  string
+}
+
+func (tm *textMap) Get(key string) string {
+	switch key {
+	case traceparentHeader:
+		return tm.parent
+	case tracestateHeader:
+		return tm.state
+	default:
+		return ""
+	}
+}
+
+func (tm *textMap) Set(key string, value string) {
+	switch key {
+	case traceparentHeader:
+		tm.parent = value
+	case tracestateHeader:
+		tm.state = value
+	}
+}
+
+func (tm *textMap) Keys() []string {
+	var k []string
+	if tm.parent != "" {
+		k = append(k, traceparentHeader)
+	}
+	if tm.state != "" {
+		k = append(k, tracestateHeader)
+	}
+	return k
+}