Datadog APMのtrace contextをmedia playlistに埋める

datadog.go

package datadog

import (
	"bytes"
	"encoding/base64"
	"encoding/binary"
	"fmt"
	"strings"
	"time"

	"github.com/grafov/m3u8"
	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace"
	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
)

const tagName = "#EXT-X-DATERANGE"
const idPrefix = "trace-"
const className = "trace-context"
const attributePrefix = "X-TRACE-"

type HLSSpanContext struct {
	SpanContext ddtrace.SpanContext
	StartDate   time.Time
	attributes  tracer.TextMapCarrier
}

func NewHLSSpanContext(ctx ddtrace.SpanContext, startDate time.Time) (*HLSSpanContext, error) {
	attributes := tracer.TextMapCarrier(map[string]string{})
	if err := tracer.Inject(ctx, attributes); err != nil {
		return nil, err
	}

	return &HLSSpanContext{
		SpanContext: ctx,
		StartDate:   startDate,
		attributes:  attributes,
	}, nil
}

// m3u8.CustomTag
func (h *HLSSpanContext) TagName() string {
	return tagName
}

// m3u8.CustomTag
func (h *HLSSpanContext) Encode() *bytes.Buffer {
	var b bytes.Buffer
	b.Write([]byte(h.String()))
	return &b
}

// m3u8.CustomTag
func (h *HLSSpanContext) String() string {
	b := make([]byte, 0, 16)
	b = binary.LittleEndian.AppendUint64(b, h.SpanContext.TraceID())
	b = binary.LittleEndian.AppendUint64(b, h.SpanContext.SpanID())
	id := idPrefix + base64.RawStdEncoding.EncodeToString(b)
	class := className
	startDate := h.StartDate.Format(time.RFC3339)

	attributes := make([]string, 0, 3+len(h.attributes))
	attributes = append(
		attributes,
		"ID="+"\""+id+"\"",
		"CLASS="+"\""+class+"\"",
		"START-DATE="+"\""+startDate+"\"",
	)
	for k, v := range h.attributes {
		attributes = append(attributes, attributePrefix+k+"="+"\""+v+"\"")
	}

	return tagName + ":" + strings.Join(attributes, ",")
}

type HLSSpanDecoder struct {
}

// m3u8.CustomDecoder
func (d *HLSSpanDecoder) TagName() string {
	return tagName
}

// m3u8.CustomDecoder
func (d *HLSSpanDecoder) Decode(line string) (m3u8.CustomTag, error) {
	if !strings.HasPrefix(line, tagName+":") {
		return nil, fmt.Errorf("invalid tag name: %s", line)
	}

	kv, err := parseAttributes(line)
	if err != nil {
		return nil, err
	}

	var class string
	var startDate time.Time
	textMap := tracer.TextMapCarrier(make(map[string]string))
	for k, v := range kv {
		switch k {
		case "CLASS":
			class = v
		case "START-DATE":
			t, err := time.Parse(time.RFC3339, v)
			if err != nil {
				return nil, err
			}
			startDate = t
		default:
			if strings.HasPrefix(k, attributePrefix) {
				k = strings.TrimPrefix(k, attributePrefix)
				textMap[k] = v
			}
		}
	}

	if class != className {
		return nil, fmt.Errorf("invalid class: %s", class)
	}

	spanContext, err := tracer.Extract(textMap)
	if err != nil {
		return nil, err
	}

	h, err := NewHLSSpanContext(spanContext, startDate)
	if err != nil {
		return nil, err
	}

	return h, nil
}

// m3u8.CustomDecoder
func (d *HLSSpanDecoder) SegmentTag() bool {
	return true
}

func parseAttributes(line string) (map[string]string, error) {
	line = strings.TrimPrefix(line, tagName+":")
	runes := []rune(line)

	attributes := make([]string, 0, 4)
	quoted := false
	split := false
	start := 0
	for i, r := range runes {
		switch {
		case split:
			split = false
			attributes = append(attributes, string(runes[start:i-1]))
			start = i
		case r == '"':
			quoted = !quoted
			continue
		case r == ',' && !quoted:
			split = true
		}
	}
	attributes = append(attributes, string(runes[start:]))

	textMap := make(map[string]string)
	for _, attribute := range attributes {
		kv := strings.Split(attribute, "=")
		if len(kv) < 2 {
			return nil, fmt.Errorf("invalid attribute: %s", attribute)
		}
		v := strings.Join(kv[1:], "=")
		v = strings.TrimPrefix(v, "\"")
		v = strings.TrimSuffix(v, "\"")
		textMap[kv[0]] = v
	}
	return textMap, nil
}

datadog_test.go

package datadog

import (
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/mocktracer"
	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
)

func TestEncodeHLSSpanContext(t *testing.T) {
	mt := mocktracer.Start()
	defer mt.Stop()

	ts, err := time.Parse(time.RFC3339, "2024-01-02T15:00:01Z")
	assert.NoError(t, err)

	span := tracer.StartSpan("test-span", tracer.WithSpanID(100))
	assert.NotNil(t, span.Context())

	h, err := NewHLSSpanContext(span.Context(), ts)
	assert.NoError(t, err)
	tag := h.String()
	assert.NotNil(t, tag)
	assert.Contains(t, []string{
		`#EXT-X-DATERANGE:ID="trace-ZAAAAAAAAABkAAAAAAAAAA",CLASS="trace-context",START-DATE="2024-01-02T15:00:01Z",X-TRACE-x-datadog-trace-id="100",X-TRACE-x-datadog-parent-id="100"`,
		`#EXT-X-DATERANGE:ID="trace-ZAAAAAAAAABkAAAAAAAAAA",CLASS="trace-context",START-DATE="2024-01-02T15:00:01Z",X-TRACE-x-datadog-parent-id="100",X-TRACE-x-datadog-trace-id="100"`,
	}, tag)
}

func TestDecodeHLSSpanContext(t *testing.T) {
	mt := mocktracer.Start()
	defer mt.Stop()

	m3u8 := `#EXT-X-DATERANGE:ID="trace-vds3qcWHqGPD4B6WKOv7IA",CLASS="trace-context",START-DATE="2024-10-28T16:17:37+09:00",X-TRACE-x-datadog-sampling-priority="1",X-TRACE-x-datadog-tags="_dd.p.dm=-1,_dd.p.tid=671f3a9100000000",X-TRACE-traceparent="00-671f3a910000000063a887c5a937dbbd-20fbeb28961ee0c3-01",X-TRACE-tracestate="dd=s:1;p:20fbeb28961ee0c3;t.dm:-1;t.tid:671f3a9100000000",X-TRACE-x-datadog-trace-id="7181138888859573181",X-TRACE-x-datadog-parent-id="2376751787917893827"`

	decoder := &HLSSpanDecoder{}
	c, err := decoder.Decode(m3u8)
	assert.NoError(t, err)
	assert.NotNil(t, c)

	sc, ok := c.(*HLSSpanContext)
	assert.True(t, ok)
	assert.NotNil(t, sc)
	assert.NotNil(t, sc.SpanContext)
	assert.Equal(t, uint64(7181138888859573181), sc.SpanContext.TraceID())
	assert.Equal(t, uint64(2376751787917893827), sc.SpanContext.SpanID())
}

func TestParseAttributes(t *testing.T) {
	t.Parallel()

	m3u8 := `#EXT-X-DATERANGE:ID="trace-vds3qcWHqGPD4B6WKOv7IA",CLASS="trace-context",START-DATE="2024-10-28T16:17:37+09:00",X-TRACE-x-datadog-sampling-priority="1",X-TRACE-x-datadog-tags="_dd.p.dm=-1,_dd.p.tid=671f3a9100000000",X-TRACE-traceparent="00-671f3a910000000063a887c5a937dbbd-20fbeb28961ee0c3-01",X-TRACE-tracestate="dd=s:1;p:20fbeb28961ee0c3;t.dm:-1;t.tid:671f3a9100000000",X-TRACE-x-datadog-trace-id="7181138888859573181",X-TRACE-x-datadog-parent-id="2376751787917893827"`
	a, err := parseAttributes(m3u8)
	assert.NoError(t, err)
	assert.NotNil(t, a)
}