switch to use spanContext object and add unmarshal code

Author: jackgopack4

exporter/exporterhelper/traces.go

@@ -5,6 +5,7 @@ package exporterhelper // import "go.opentelemetry.io/collector/exporter/exporte
 
 import (
 	"context"
+	"encoding/hex"
 	"encoding/json"
 	"errors"
 
@@ -46,55 +47,6 @@ func NewTracesQueueBatchSettings() QueueBatchSettings {
 	}
 }
 
-type SerializableLink struct {
-	TraceID    [16]byte         `json:"trace_id"`
-	SpanID     [8]byte          `json:"span_id"`
-	TraceFlags byte             `json:"trace_flags"`
-	TraceState trace.TraceState `json:"trace_state"`
-}
-
-func (sl *SerializableLink) UnmarshalJSON(data []byte) error {
-	type Alias SerializableLink // Prevent recursion
-	aux := &struct {
-		TraceState string `json:"trace_state"`
-		*Alias
-	}{
-		Alias: (*Alias)(sl),
-	}
-	if err := json.Unmarshal(data, &aux); err != nil {
-		return err
-	}
-	if aux.TraceState != "" {
-		ts, err := trace.ParseTraceState(aux.TraceState)
-		if err != nil {
-			return err
-		}
-		sl.TraceState = ts
-	}
-	return nil
-}
-
-func linkToSerializable(l trace.Link) SerializableLink {
-	return SerializableLink{
-		TraceID:    l.SpanContext.TraceID(),
-		SpanID:     l.SpanContext.SpanID(),
-		TraceFlags: byte(l.SpanContext.TraceFlags()),
-		TraceState: l.SpanContext.TraceState(),
-	}
-}
-
-func serializableToLink(sl SerializableLink) trace.Link {
-	sc := trace.NewSpanContext(trace.SpanContextConfig{
-		TraceID:    sl.TraceID,
-		SpanID:     sl.SpanID,
-		TraceFlags: trace.TraceFlags(sl.TraceFlags),
-		TraceState: sl.TraceState,
-	})
-	return trace.Link{
-		SpanContext: sc,
-	}
-}
-
 type tracesRequest struct {
 	td         ptrace.Traces
 	links      []trace.Link
@@ -112,22 +64,72 @@ func newTracesRequest(td ptrace.Traces, links []trace.Link) Request {
 type tracesEncoding struct{}
 
 type tracesWithSpanContexts struct {
-	Traces []byte             `json:"traces"`
-	Links  []SerializableLink `json:"links"`
+	Traces      []byte              `json:"traces"`
+	SpanContext []trace.SpanContext `json:"span_context"`
+}
+
+// Helper for JSON unmarshaling of SpanContextConfig
+
+type spanContextConfigJSON struct {
+	TraceID    string
+	SpanID     string
+	TraceFlags string
+	TraceState string
+	Remote     bool
+}
+
+func unmarshalSpanContextConfig(data []byte) (trace.SpanContextConfig, error) {
+	var aux spanContextConfigJSON
+	if err := json.Unmarshal(data, &aux); err != nil {
+		return trace.SpanContextConfig{}, err
+	}
+	tid, err := trace.TraceIDFromHex(aux.TraceID)
+	if err != nil {
+		return trace.SpanContextConfig{}, err
+	}
+	sid, err := trace.SpanIDFromHex(aux.SpanID)
+	if err != nil {
+		return trace.SpanContextConfig{}, err
+	}
+	ts, err := trace.ParseTraceState(aux.TraceState)
+	if err != nil {
+		return trace.SpanContextConfig{}, err
+	}
+	tf, err := hex.DecodeString(aux.TraceFlags)
+	if err != nil {
+		return trace.SpanContextConfig{}, err
+	}
+	if len(tf) != 1 {
+		return trace.SpanContextConfig{}, errors.New("invalid trace flags")
+	}
+	return trace.SpanContextConfig{
+		TraceID:    tid,
+		SpanID:     sid,
+		TraceFlags: trace.TraceFlags(tf[0]),
+		TraceState: ts,
+		Remote:     aux.Remote,
+	}, nil
 }
 
 func (tracesEncoding) Unmarshal(bytes []byte) (Request, error) {
-	var twl tracesWithSpanContexts
+	var twl struct {
+		Traces      []byte            `json:"traces"`
+		SpanContext []json.RawMessage `json:"span_context"`
+	}
 	if err := json.Unmarshal(bytes, &twl); err != nil {
 		return nil, err
 	}
 	traces, err := tracesUnmarshaler.UnmarshalTraces(twl.Traces)
 	if err != nil {
 		return nil, err
 	}
-	links := make([]trace.Link, len(twl.Links))
-	for i, sl := range twl.Links {
-		links[i] = serializableToLink(sl)
+	links := make([]trace.Link, len(twl.SpanContext))
+	for i, raw := range twl.SpanContext {
+		cfg, err := unmarshalSpanContextConfig(raw)
+		if err != nil {
+			return nil, err
+		}
+		links[i] = trace.Link{SpanContext: trace.NewSpanContext(cfg)}
 	}
 	return newTracesRequest(traces, links), nil
 }
@@ -138,13 +140,13 @@ func (tracesEncoding) Marshal(req Request) ([]byte, error) {
 	if err != nil {
 		return nil, err
 	}
-	serializableLinks := make([]SerializableLink, len(tr.links))
+	spanContexts := make([]trace.SpanContext, len(tr.links))
 	for i, l := range tr.links {
-		serializableLinks[i] = linkToSerializable(l)
+		spanContexts[i] = l.SpanContext
 	}
 	twl := tracesWithSpanContexts{
-		Traces: tracesBytes,
-		Links:  serializableLinks,
+		Traces:      tracesBytes,
+		SpanContext: spanContexts,
 	}
 	return json.Marshal(twl)
 }

exporter/exporterhelper/traces_test.go

@@ -5,6 +5,7 @@ package exporterhelper
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
 	"testing"
 	"time"
@@ -480,18 +481,27 @@ func checkWrapSpanForTraces(t *testing.T, sr *tracetest.SpanRecorder, tracer tra
 
 func TestSerializableToLink(t *testing.T) {
 	ts, err := trace.TraceState{}.Insert("key", "value")
-	if err != nil {
-		t.Fatalf("Error inserting trace state: %v", err)
-	}
-	sl := SerializableLink{
+	require.NoError(t, err)
+
+	sc := trace.NewSpanContext(trace.SpanContextConfig{
 		TraceID:    [16]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16},
 		SpanID:     [8]byte{1, 2, 3, 4, 5, 6, 7, 8},
-		TraceFlags: byte(trace.FlagsSampled),
+		TraceFlags: trace.FlagsSampled,
 		TraceState: ts,
-	}
-	res := serializableToLink(sl)
-	// TraceState will be empty due to error in inserting Member Key "@key"
-	assert.Equal(t, ts, res.SpanContext.TraceState())
+	})
+
+	// Marshal and unmarshal SpanContext as JSON
+	data, err := json.Marshal(sc)
+	require.NoError(t, err)
+
+	sc2, err := unmarshalSpanContextConfig(data)
+	require.NoError(t, err)
+	scRestored := trace.NewSpanContext(sc2)
+
+	assert.Equal(t, sc.TraceID(), scRestored.TraceID())
+	assert.Equal(t, sc.SpanID(), scRestored.SpanID())
+	assert.Equal(t, sc.TraceFlags(), scRestored.TraceFlags())
+	assert.Equal(t, sc.TraceState(), scRestored.TraceState())
 }
 
 func TestTracesEncoding_Unmarshal_InvalidJSON(t *testing.T) {