WithTraceContext replaces context instead of appending

envelope/tracecontext.go

@@ -25,9 +25,13 @@ func WithTraceContext(ctx context.Context, e messenger.Envelope) messenger.Envel
 		return e
 	}
 
+	e = envelope.WithoutHeader(e, TraceparentHeader)
+	e = envelope.WithoutHeader(e, TracestateHeader)
+
 	tc := tx.TraceContext()
 	traceparentValue := apmhttp.FormatTraceparentHeader(tc)
 	tracestateValue := tc.State.String()
+
 	e = envelope.WithHeader(e, TraceparentHeader, traceparentValue)
 	e = envelope.WithHeader(e, TracestateHeader, tracestateValue)
 
@@ -36,11 +40,6 @@ func WithTraceContext(ctx context.Context, e messenger.Envelope) messenger.Envel
 
 // StartTransaction start transaction using W3C Trace-Context headers in the envelope
 func StartTransaction(tracer *apm.Tracer, name string, typ string, e messenger.Envelope) *apm.Transaction {
-	messageType := envelope.MessageType(e)
-	if len(messageType) == 0 {
-		messageType = "unknown"
-	}
-
 	traceparentValue, _ := e.LastHeader(TraceparentHeader)
 	traceContext, _ := apmhttp.ParseTraceparentHeader(traceparentValue)
 

envelope/tracecontext_test.go

@@ -11,19 +11,22 @@ import (
 	"testing"
 )
 
-func TestWithTraceContext_given_context_with_transaction_it_should_add_trace_context_headers_with_its_value(t *testing.T) {
+func TestWithTraceContext_given_context_with_transaction_it_should_set_trace_context_headers_with_its_value(t *testing.T) {
 	tracer := apmtest.NewDiscardTracer()
-	tx := tracer.StartTransaction("Test", "test")
+	tx1 := tracer.StartTransaction("Test1", "test")
+	tx2 := tracer.StartTransaction("Test2", "test")
 
-	ctx := apm.ContextWithTransaction(context.Background(), tx)
+	ctx1 := apm.ContextWithTransaction(context.Background(), tx1)
+	ctx2 := apm.ContextWithTransaction(context.Background(), tx2)
 
-	e := WithTraceContext(ctx, envelope.FromMessage("test message"))
+	e := WithTraceContext(ctx1, envelope.FromMessage("test message"))
+	e = WithTraceContext(ctx2, e)
 
-	traceparent, _ := e.LastHeader(TraceparentHeader)
-	tracestate, _ := e.LastHeader(TracestateHeader)
+	traceparent := e.Header(TraceparentHeader)
+	tracestate := e.Header(TracestateHeader)
 
-	assert.Equal(t, traceparent, apmhttp.FormatTraceparentHeader(tx.TraceContext()))
-	assert.Equal(t, tracestate, tx.TraceContext().State.String())
+	assert.Equal(t, traceparent, []string{apmhttp.FormatTraceparentHeader(tx2.TraceContext())})
+	assert.Equal(t, tracestate, []string{tx2.TraceContext().State.String()})
 }
 
 func TestWithTraceContext_given_context_without_transaction_it_should_return_same_envelope_without_change(t *testing.T) {

go.mod

@@ -1,11 +1,36 @@
 module github.com/riid/messenger-apm
 
-go 1.16
+go 1.17
 
 require (
-	github.com/riid/messenger v0.0.0-20220114175347-af9da6925871
+	github.com/riid/messenger v0.0.0-20220117162302-234459211f7b
 	go.elastic.co/apm v1.15.0
 	go.elastic.co/apm/module/apmhttp v1.15.0
 )
 
 require github.com/stretchr/testify v1.7.0
+
+require (
+	github.com/armon/go-radix v1.0.0 // indirect
+	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/elastic/go-licenser v0.3.1 // indirect
+	github.com/elastic/go-sysinfo v1.1.1 // indirect
+	github.com/elastic/go-windows v1.0.0 // indirect
+	github.com/google/go-cmp v0.3.1 // indirect
+	github.com/jcchavezs/porto v0.1.0 // indirect
+	github.com/joeshaw/multierror v0.0.0-20140124173710-69b34d4ec901 // indirect
+	github.com/kr/pretty v0.2.1 // indirect
+	github.com/pkg/errors v0.8.1 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
+	github.com/prometheus/procfs v0.0.0-20190425082905-87a4384529e0 // indirect
+	github.com/santhosh-tekuri/jsonschema v1.2.4 // indirect
+	github.com/stretchr/objx v0.1.1 // indirect
+	go.elastic.co/fastjson v1.1.0 // indirect
+	golang.org/x/lint v0.0.0-20201208152925-83fdc39ff7b5 // indirect
+	golang.org/x/mod v0.3.0 // indirect
+	golang.org/x/sys v0.0.0-20191204072324-ce4227a45e2e // indirect
+	golang.org/x/tools v0.0.0-20200509030707-2212a7e161a5 // indirect
+	golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 // indirect
+	gopkg.in/yaml.v3 v3.0.0-20200615113413-eeeca48fe776 // indirect
+	howett.net/plist v0.0.0-20181124034731-591f970eefbb // indirect
+)

go.sum

@@ -29,8 +29,8 @@ github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZb
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/prometheus/procfs v0.0.0-20190425082905-87a4384529e0 h1:c8R11WC8m7KNMkTv/0+Be8vvwo4I3/Ut9AC2FW8fX3U=
 github.com/prometheus/procfs v0.0.0-20190425082905-87a4384529e0/go.mod h1:TjEm7ze935MbeOT/UhFTIMYKhuLP4wbCsTZCD3I8kEA=
-github.com/riid/messenger v0.0.0-20220114175347-af9da6925871 h1:Zq8O/OA3qGtX5Bp7XAst6kQvWsi+dnX21QDJotrP//0=
-github.com/riid/messenger v0.0.0-20220114175347-af9da6925871/go.mod h1:jt0PPNl3y7L7JWfs1fk+jwHdPaB2C29XlfajdTLmU+0=
+github.com/riid/messenger v0.0.0-20220117162302-234459211f7b h1:vwQTLoTgcZlt3ttMk3skZ4Ect3UZd7jqErtKdRDzUgo=
+github.com/riid/messenger v0.0.0-20220117162302-234459211f7b/go.mod h1:jt0PPNl3y7L7JWfs1fk+jwHdPaB2C29XlfajdTLmU+0=
 github.com/santhosh-tekuri/jsonschema v1.2.4 h1:hNhW8e7t+H1vgY+1QeEQpveR6D4+OwKPXCfD2aieJis=
 github.com/santhosh-tekuri/jsonschema v1.2.4/go.mod h1:TEAUOeZSmIxTTuHatJzrvARHiuO9LYd+cIxzgEHCQI4=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=