Repository: elastic/apm-agent-go (go)

Issue:
Can't mix native transaction with Otel's `trace.SpanFromContext`
**Describe the bug**

The OTel bridge have the same problem that the opentracing bridge once faced (#235 and https://github.com/opentracing/opentracing-go/issues/192)

apm’s native transaction can't be retrieved by Otel’s `trace.SpanFromContext`(https://github.com/open-telemetry/opentelemetry-go/blob/b2246d58650bc6db04459353843250ee3bba77fc/trace/context.go#L48).

Since goredis (https://github.com/redis/go-redis/blob/dc8c93ba66777b8f9866542b47cd102f591b5984/extra/redisotel/tracing.go#L90) and maybe other libs use `trace.SpanFromContext` before starting a Span,
currently it is not possible to mix native apm instrumentation with those lib’s OTel instrumentation.

**To Reproduce**

```go
package main

import (
	"context"
	"fmt"

	"go.elastic.co/apm/module/apmotel/v2"
	"go.elastic.co/apm/v2"
	"go.elastic.co/apm/v2/apmtest"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/trace"
)

func main() {
	tracer := apmtest.NewRecordingTracer()
	provider, err := apmotel.NewTracerProvider(apmotel.WithAPMTracer(tracer.Tracer))
	if err != nil {
		panic(err)
	}
	otel.SetTracerProvider(provider)

	ctx := context.Background()
	tx := tracer.StartTransaction("test1", "test")
	ctx = apm.ContextWithTransaction(ctx, tx)

	fmt.Printf("apm transaction: %#v\n", apm.TransactionFromContext(ctx))
	sp := trace.SpanFromContext(ctx)
	fmt.Printf("OTel span: %#v\n", sp) // expect to not be otel trace.noopSpan{}
}

```

**Expected behavior**

Should be able to retrieve Span using Otel’s `trace.SpanFromContext` from an apm transaction.


Pull request:
Allow retrieving span created in agent from otel, and store agent span in context when creating it with otel
Closes #1449 

We didn't store the agent span/transaction into the context, as that doesn't happen automatically and requires a manual operation when using the agent directly.
When using otel, which does add it to every context, it makes more sense not to have to require folks to do that themselves though.

This also sets up a context wrapper in the agent, so spans created from the agent are stored in the context to be read by otel.

Patch:
diff --git a/module/apmotel/tracer.go b/module/apmotel/tracer.go
index e3ec0672..2205688f 100644
--- a/module/apmotel/tracer.go
+++ b/module/apmotel/tracer.go
@@ -86,6 +86,7 @@ func (t *tracer) Start(ctx context.Context, spanName string, opts ...trace.SpanS
 				Parent: tc,
 				Start:  startTime,
 			})
+			ctx = apm.ContextWithSpan(ctx, s.span)
 			s.tx = parent.tx
 			s.spanContext = trace.NewSpanContext(trace.SpanContextConfig{
 				TraceID:    trace.TraceID(s.span.TraceContext().Trace),
@@ -105,6 +106,7 @@ func (t *tracer) Start(ctx context.Context, spanName string, opts ...trace.SpanS
 		}
 	}
 	s.tx = t.provider.apmTracer.StartTransactionOptions(spanName, "", tranOpts)
+	ctx = apm.ContextWithTransaction(ctx, s.tx)
 	s.spanContext = trace.NewSpanContext(trace.SpanContextConfig{
 		TraceID:    trace.TraceID(s.tx.TraceContext().Trace),
 		SpanID:     trace.SpanID(s.tx.TraceContext().Span),
diff --git a/module/apmotel/wrapper.go b/module/apmotel/wrapper.go
new file mode 100644
index 00000000..f3bffc2d
--- /dev/null
+++ b/module/apmotel/wrapper.go
@@ -0,0 +1,63 @@
+// Licensed to Elasticsearch B.V. under one or more contributor
+// license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright
+// ownership. Elasticsearch B.V. licenses this file to you under
+// the Apache License, Version 2.0 (the "License"); you may
+// not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//go:build go1.18
+// +build go1.18
+
+package apmotel // import "go.elastic.co/apm/module/apmotel/v2"
+
+import (
+	"context"
+
+	"go.opentelemetry.io/otel/trace"
+
+	"go.elastic.co/apm/v2"
+)
+
+var (
+	oldOverrideContextWithSpan        func(context.Context, *apm.Span) context.Context
+	oldOverrideContextWithTransaction func(context.Context, *apm.Transaction) context.Context
+)
+
+func init() {
+	// We override the apm context functions so that transactions
+	// and spans started with the native API are wrapped and made
+	// available as OpenTelemetry spans.
+
+	oldOverrideContextWithSpan = apm.OverrideContextWithSpan
+	oldOverrideContextWithTransaction = apm.OverrideContextWithTransaction
+	apm.OverrideContextWithSpan = contextWithSpan
+	apm.OverrideContextWithTransaction = contextWithTransaction
+}
+
+func contextWithSpan(ctx context.Context, apmSpan *apm.Span) context.Context {
+	ctx = oldOverrideContextWithSpan(ctx, apmSpan)
+	return trace.ContextWithSpanContext(ctx, trace.NewSpanContext(trace.SpanContextConfig{
+		TraceID:    trace.TraceID(apmSpan.TraceContext().Trace),
+		SpanID:     trace.SpanID(apmSpan.TraceContext().Span),
+		TraceFlags: trace.TraceFlags(0).WithSampled(!apmSpan.Dropped()),
+	}))
+}
+
+func contextWithTransaction(ctx context.Context, apmTransaction *apm.Transaction) context.Context {
+	ctx = oldOverrideContextWithTransaction(ctx, apmTransaction)
+	return trace.ContextWithSpanContext(ctx, trace.NewSpanContext(trace.SpanContextConfig{
+		TraceID:    trace.TraceID(apmTransaction.TraceContext().Trace),
+		SpanID:     trace.SpanID(apmTransaction.TraceContext().Span),
+		TraceFlags: trace.TraceFlags(0).WithSampled(apmTransaction.Sampled()),
+	}))
+}


Test patch:
diff --git a/module/apmotel/tracer_test.go b/module/apmotel/tracer_test.go
index 811fccbb..9a4a717b 100644
--- a/module/apmotel/tracer_test.go
+++ b/module/apmotel/tracer_test.go
@@ -56,6 +56,9 @@ func TestTracerStartTransaction(t *testing.T) {
 
 	assert.True(t, trace.SpanContextFromContext(ctx).IsValid())
 	assert.True(t, trace.SpanContextFromContext(ctx).IsSampled())
+
+	assert.True(t, apm.TransactionFromContext(ctx).Sampled())
+	assert.Nil(t, apm.SpanFromContext(ctx))
 }
 
 func TestTracerStartTransactionWithParentContext(t *testing.T) {
@@ -124,6 +127,9 @@ func TestTracerStartChildSpan(t *testing.T) {
 	assert.NotNil(t, cs.(*span).span)
 
 	assert.True(t, trace.SpanContextFromContext(ctx).IsValid())
+
+	assert.True(t, apm.TransactionFromContext(ctx).Sampled())
+	assert.False(t, apm.SpanFromContext(ctx).Dropped())
 }
 
 func TestTracerStartChildSpanFromTransactionInContext(t *testing.T) {
@@ -162,7 +168,6 @@ func TestTracerStartChildSpanFromTransactionInContext(t *testing.T) {
 		},
 	} {
 		t.Run(tt.name, func(t *testing.T) {
-
 			apmTracer, _ := transporttest.NewRecorderTracer()
 			tp, err := NewTracerProvider(WithAPMTracer(apmTracer))
 			assert.NoError(t, err)
diff --git a/module/apmotel/wrapper_test.go b/module/apmotel/wrapper_test.go
new file mode 100644
index 00000000..b9146a97
--- /dev/null
+++ b/module/apmotel/wrapper_test.go
@@ -0,0 +1,63 @@
+// Licensed to Elasticsearch B.V. under one or more contributor
+// license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright
+// ownership. Elasticsearch B.V. licenses this file to you under
+// the Apache License, Version 2.0 (the "License"); you may
+// not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//go:build go1.18
+// +build go1.18
+
+package apmotel
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"go.opentelemetry.io/otel/trace"
+
+	"go.elastic.co/apm/v2"
+	"go.elastic.co/apm/v2/transport/transporttest"
+)
+
+func TestLinkAgentToOtel(t *testing.T) {
+	apmTracer, _ := transporttest.NewRecorderTracer()
+	_, err := NewTracerProvider(WithAPMTracer(apmTracer))
+	assert.NoError(t, err)
+
+	ctx := context.Background()
+	tx := apmTracer.StartTransaction("test1", "test")
+	ctx = apm.ContextWithTransaction(ctx, tx)
+
+	apmTx := apm.TransactionFromContext(ctx)
+	otelSpan := trace.SpanFromContext(ctx)
+
+	assert.Equal(t, [16]byte(apmTx.TraceContext().Trace), [16]byte(otelSpan.SpanContext().TraceID()))
+	assert.Equal(t, [8]byte(apmTx.TraceContext().Span), [8]byte(otelSpan.SpanContext().SpanID()))
+}
+
+func TestLinkOtelToAgent(t *testing.T) {
+	apmTracer, _ := transporttest.NewRecorderTracer()
+	tp, err := NewTracerProvider(WithAPMTracer(apmTracer))
+	assert.NoError(t, err)
+
+	ctx := context.Background()
+	ctx, _ = tp.Tracer("").Start(ctx, "test")
+
+	apmTx := apm.TransactionFromContext(ctx)
+	otelSpan := trace.SpanFromContext(ctx)
+
+	assert.Equal(t, [16]byte(apmTx.TraceContext().Trace), [16]byte(otelSpan.SpanContext().TraceID()))
+	assert.Equal(t, [8]byte(apmTx.TraceContext().Span), [8]byte(otelSpan.SpanContext().SpanID()))
+}