replace serializableTraceState with tracestate

Author: jackgopack4

exporter/exporterhelper/traces.go

@@ -46,81 +46,49 @@ func NewTracesQueueBatchSettings() QueueBatchSettings {
 	}
 }
 
-// TraceStateSerializable represents a serializable version of TraceState
-type TraceStateSerializable struct {
-	Members []struct {
-		Key   string `json:"key"`
-		Value string `json:"value"`
-	} `json:"members"`
+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"`
 }
 
-// traceStateToSerializable converts a TraceState to a serializable format
-func traceStateToSerializable(ts trace.TraceState) TraceStateSerializable {
-	result := TraceStateSerializable{
-		Members: make([]struct {
-			Key   string `json:"key"`
-			Value string `json:"value"`
-		}, 0, ts.Len()),
+func (sl *SerializableLink) UnmarshalJSON(data []byte) error {
+	type Alias SerializableLink // Prevent recursion
+	aux := &struct {
+		TraceState string `json:"trace_state"`
+		*Alias
+	}{
+		Alias: (*Alias)(sl),
 	}
-
-	ts.Walk(func(key, value string) bool {
-		result.Members = append(result.Members, struct {
-			Key   string `json:"key"`
-			Value string `json:"value"`
-		}{
-			Key:   key,
-			Value: value,
-		})
-		return true
-	})
-
-	return result
-}
-
-// serializableToTraceState converts a serializable format back to TraceState
-func serializableToTraceState(ts TraceStateSerializable) (trace.TraceState, error) {
-	var result trace.TraceState
-	var err error
-
-	for _, m := range ts.Members {
-		result, err = result.Insert(m.Key, m.Value)
+	if err := json.Unmarshal(data, &aux); err != nil {
+		return err
+	}
+	if aux.TraceState != "" {
+		ts, err := trace.ParseTraceState(aux.TraceState)
 		if err != nil {
-			return trace.TraceState{}, err
+			return err
 		}
+		sl.TraceState = ts
 	}
-
-	return result, nil
-}
-
-// Update SerializableLink to use the new TraceState serialization
-type SerializableLink struct {
-	TraceID    [16]byte               `json:"trace_id"`
-	SpanID     [8]byte                `json:"span_id"`
-	TraceFlags byte                   `json:"trace_flags"`
-	TraceState TraceStateSerializable `json:"trace_state"`
+	return nil
 }
 
 func linkToSerializable(l trace.Link) SerializableLink {
 	return SerializableLink{
 		TraceID:    l.SpanContext.TraceID(),
 		SpanID:     l.SpanContext.SpanID(),
 		TraceFlags: byte(l.SpanContext.TraceFlags()),
-		TraceState: traceStateToSerializable(l.SpanContext.TraceState()),
+		TraceState: l.SpanContext.TraceState(),
 	}
 }
 
 func serializableToLink(sl SerializableLink) trace.Link {
-	ts, err := serializableToTraceState(sl.TraceState)
-	if err != nil {
-		// If there's an error parsing the trace state, use an empty one
-		ts = trace.TraceState{}
-	}
-
 	sc := trace.NewSpanContext(trace.SpanContextConfig{
 		TraceID:    sl.TraceID,
 		SpanID:     sl.SpanID,
 		TraceFlags: trace.TraceFlags(sl.TraceFlags),
-		TraceState: ts,
+		TraceState: sl.TraceState,
 	})
 	return trace.Link{
 		SpanContext: sc,
@@ -143,13 +111,13 @@ func newTracesRequest(td ptrace.Traces, links []trace.Link) Request {
 
 type tracesEncoding struct{}
 
-type tracesWithLinks struct {
+type tracesWithSpanContexts struct {
 	Traces []byte             `json:"traces"`
 	Links  []SerializableLink `json:"links"`
 }
 
 func (tracesEncoding) Unmarshal(bytes []byte) (Request, error) {
-	var twl tracesWithLinks
+	var twl tracesWithSpanContexts
 	if err := json.Unmarshal(bytes, &twl); err != nil {
 		return nil, err
 	}
@@ -174,7 +142,7 @@ func (tracesEncoding) Marshal(req Request) ([]byte, error) {
 	for i, l := range tr.links {
 		serializableLinks[i] = linkToSerializable(l)
 	}
-	twl := tracesWithLinks{
+	twl := tracesWithSpanContexts{
 		Traces: tracesBytes,
 		Links:  serializableLinks,
 	}

exporter/exporterhelper/traces_test.go

@@ -479,25 +479,19 @@ 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{
 		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),
-		TraceState: TraceStateSerializable{
-			Members: []struct {
-				Key   string `json:"key"`
-				Value string `json:"value"`
-			}{
-				{
-					Key:   "@key", // @ is not allowed in trace state keys
-					Value: "value",
-				},
-			},
-		},
+		TraceState: ts,
 	}
 	res := serializableToLink(sl)
 	// TraceState will be empty due to error in inserting Member Key "@key"
-	assert.Equal(t, trace.TraceState{}, res.SpanContext.TraceState())
+	assert.Equal(t, ts, res.SpanContext.TraceState())
 }
 
 func TestTracesEncoding_Unmarshal_InvalidJSON(t *testing.T) {