equinix-labs · tobert · Sep 22, 2021 · Sep 22, 2021 · Sep 22, 2021 · Sep 22, 2021
diff --git a/go.mod b/go.mod
@@ -7,5 +7,6 @@ require (
 	go.opentelemetry.io/otel v1.0.0-RC2
 	go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.0.0-RC2
 	go.opentelemetry.io/otel/sdk v1.0.0-RC2
+	go.opentelemetry.io/otel/trace v1.0.0-RC2
 	google.golang.org/grpc v1.39.0
 )
diff --git a/otelhelpers/context_traceparent.go b/otelhelpers/context_traceparent.go
@@ -0,0 +1,87 @@
+// otelhelpers is a package of helper functions for dealing with otel
+// traceparent propagation over files and environment variables.
+package otelhelpers
+
+import (
+	"bytes"
+	"context"
+	"io/ioutil"
+	"os"
+
+	"go.opentelemetry.io/otel"
+)
+
+// ContextWithEnvTraceparent is a helper that looks for the the TRACEPARENT
+// environment variable and if it's set, it grabs the traceparent and
+// adds it to the context it returns. When there is no envvar or it's
+// empty, the original context is returned unmodified.
+func ContextWithEnvTraceparent(ctx context.Context) context.Context {
+	traceparent := os.Getenv("TRACEPARENT")
+	if traceparent != "" {
+		return ContextWithTraceparentString(ctx, traceparent)
+	}
+	return ctx
+}
+
+// ContextWithLinuxCmdlineTraceparent looks in /proc/cmdline for a traceparent=
+// command line option and returns the context with that value as traceparent
+// if it's there. Does no validation. Returns the original context if there is
+// no cmdline option or if there's an error doing the read.
+// This is Linux-only but should be safe on other operating systems.
+func ContextWithCmdlineTraceparent(ctx context.Context) context.Context {
+	tp, err := tpFromCmdline("/proc/cmdline")
+	if err != nil {
+		// what to do with error? is there a way to hit the otel error handler infra?
+		return ctx
+	}
+
+	return ContextWithTraceparentString(ctx, tp)
+}
+
+// ContextWithCmdlineOrEnvTraceparent checks the environment variable first,
+// then /proc/cmdline and returns a context with them set, if available. When
+// both are present, the cmdline is prioritized. When neither is present,
+// the original context is returned as-is.
+func ContextWithCmdlineOrEnvTraceparent(ctx context.Context) context.Context {
+	ctx = ContextWithEnvTraceparent(ctx)
+	return ContextWithCmdlineTraceparent(ctx)
+}
+
+// ContextWithTraceparentString takes a W3C traceparent string, uses the otel
+// carrier code to get it into a context it returns ready to go.
+func ContextWithTraceparentString(ctx context.Context, traceparent string) context.Context {
+	carrier := SimpleCarrier{}
+	carrier.Set("traceparent", traceparent)
+	prop := otel.GetTextMapPropagator()
+	return prop.Extract(ctx, carrier)
+}
+
+// TraceparentStringFromContext gets the current trace from the context and
+// returns a W3C traceparent string.
+func TraceparentStringFromContext(ctx context.Context) string {
+	carrier := SimpleCarrier{}
+	prop := otel.GetTextMapPropagator()
+	prop.Inject(ctx, carrier)
+	return carrier.Get("traceparent")
+}
+
+// tpFromCmdline reads a /proc/cmdline style file, parses it, and returns whatever
+// value is present for "traceparent=".
+func tpFromCmdline(file string) (string, error) {
+	data, err := ioutil.ReadFile(file)
+	if err != nil {
+		return "", err
+	}
+
+	if bytes.Contains(data, []byte("traceparent=")) {
+		kvpairs := bytes.Split(data, []byte(" "))
+		for _, kv := range kvpairs {
+			parts := bytes.SplitN(kv, []byte("="), 2)
+			if string(parts[0]) == "traceparent" {
+				return string(parts[1]), nil
+			}
+		}
+	}
+
+	return "", nil
+}
diff --git a/otelhelpers/context_traceparent_test.go b/otelhelpers/context_traceparent_test.go
@@ -0,0 +1,76 @@
+package otelhelpers
+
+import (
+	"context"
+	"io/ioutil"
+	"os"
+	"testing"
+
+	"github.com/equinix-labs/otel-init-go/otelinit"
+	"go.opentelemetry.io/otel/trace"
+)
+
+func TestMain(m *testing.M) {
+	// Many of the tests here won't work at all if otel is in non-recording mode
+	// so we set a default endpoint and let it fail in the background.  If there
+	// happens to be a listener it will connnect but should not receive spans.
+	os.Setenv("OTEL_EXPORTER_OTLP_ENDPOINT", "localhost:4317")
+	otelinit.InitOpenTelemetry(context.Background(), "otel-init-go-helpers-test")
+
+	os.Exit(m.Run())
+}
+
+func TestContextWithEnvTraceparent(t *testing.T) {
+	// make sure the environment variable isn't polluting test state
+	os.Unsetenv("TRACEPARENT")
+
+	// trace id should not change, because there's no envvar and no file
+	ctx := ContextWithEnvTraceparent(context.Background())
+	sc := trace.SpanContextFromContext(ctx)
+	if sc.HasTraceID() {
+		t.Error("traceparent detected where there should be none")
+	}
+
+	os.Setenv("TRACEPARENT", "00-f61fc53f926e07a9c3893b1a722e1b65-7a2d6a804f3de137-01")
+	ctx = ContextWithEnvTraceparent(context.Background())
+	sc = trace.SpanContextFromContext(ctx)
+	if sc.TraceID().String() != "f61fc53f926e07a9c3893b1a722e1b65" {
+		t.Errorf("no trace id where one is expected. got: %q", sc.TraceID().String())
+	}
+	if sc.SpanID().String() != "7a2d6a804f3de137" {
+		t.Errorf("no span id where one is expected. got: %q", sc.SpanID().String())
+	}
+	if !sc.IsSampled() {
+		t.Error("expected sampling to be enabled but it is not")
+	}
+}
+
+func TestTpFromCmdline(t *testing.T) {
+	testTp := "00-f61fc53f926e07a9c3893b1a722e1b65-7a2d6a804f3de137-01"
+	testCmdlines := []string{
+		"traceparent=00-f61fc53f926e07a9c3893b1a722e1b65-7a2d6a804f3de137-01",
+		"foo=bar initrd=lol root=wheeeeeeeeeeee-fun traceparent=00-f61fc53f926e07a9c3893b1a722e1b65-7a2d6a804f3de137-01",
+		"traceparent=00-f61fc53f926e07a9c3893b1a722e1b65-7a2d6a804f3de137-01 foo=bar initrd=lol root=wheeeeeeeeeeee-fun",
+		"kimi=ga baka=desu traceparent=00-f61fc53f926e07a9c3893b1a722e1b65-7a2d6a804f3de137-01 foo=bar initrd=lol root=wheeeeeeeeeeee-fun",
+	}
+
+	for _, cmdline := range testCmdlines {
+		file, err := ioutil.TempFile(t.TempDir(), "go-test-otel-init-go")
+		if err != nil {
+			t.Fatalf("unable to create tempfile for testing: %s", err)
+		}
+		defer os.Remove(file.Name())
+
+		// write out a cmdline file for test
+		file.WriteString(cmdline)
+		file.Close()
+
+		got, err := tpFromCmdline(file.Name())
+		if err != nil {
+			t.Errorf("reading cmdline test file failed unexpectedly: %s", err)
+		}
+		if got != testTp {
+			t.Errorf("tpFromCmdline comparison failed, expected '%s', got '%s'", testTp, got)
+		}
+	}
+}
diff --git a/otelhelpers/env_traceparent.go b/otelhelpers/env_traceparent.go
diff --git a/otelhelpers/simple_carrier_test.go b/otelhelpers/simple_carrier_test.go
@@ -0,0 +1,59 @@
+package otelhelpers
+
+import (
+	"context"
+	"testing"
+
+	"go.opentelemetry.io/otel"
+)
+
+func TestSimpleCarrier(t *testing.T) {
+	carrier := SimpleCarrier{}
+	carrier.Clear() // clean up after other tests
+
+	// traceparent is the only key supported by SimpleCarrier
+	got := carrier.Get("traceparent")
+	if got != "" {
+		t.Errorf("got a non-empty traceparent value '%s' where empty string was expected", got)
+	}
+
+	carrier.Set("foobar", "baz")
+	if carrier.Get("foobar") != "baz" {
+		t.Error("did not get the expected value back in Set/Get test")
+	}
+
+	// traceparent is supported so this should work fine
+	tp := "00-b122b620341449410b9cd900c96d459d-aa21cda35388b694-01"
+	carrier.Set("traceparent", tp)
+
+	// we've set 2 keys so far, and both should get returned
+	keys := carrier.Keys()
+	if len(keys) != 2 {
+		t.Errorf("expected exactly 2 keys from Keys() but instead got %q", keys)
+	}
+
+	// make sure the value round-trips in one piece
+	got = carrier.Get("traceparent")
+	if got != tp {
+		t.Errorf("expected traceparent value '%s' but got '%s'", tp, got)
+	}
+
+	// it's impractical to test the internal state of otel-go, so the next best
+	// thing is to round-trip our traceparent through it and make sure it comes
+	// back as expected
+	prop := otel.GetTextMapPropagator()
+	ctx := prop.Extract(context.Background(), carrier)
+	if ctx == nil {
+		t.Errorf("expected a context but got nil, likely a problem in otel? this shouldn't happen...")
+	}
+
+	// try to round trip the traceparent back out of that context ^^
+	rtCarrier := SimpleCarrier{}
+	prop.Inject(ctx, rtCarrier)
+	got = carrier.Get("traceparent")
+	if got != tp {
+		t.Errorf("round-tripping traceparent through a context failed, expected '%s', got '%s'", tp, got)
+	}
+
+	carrier.Clear() // clean up for other tests
+}