docdb/internal/util/observability/propagation.go

// Copyright 2021 Hanzo AI Inc.
//
// Licensed 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.

package observability

import (
	"encoding/json"

	"github.com/AlekSi/lazyerrors"
	oteltrace "go.opentelemetry.io/otel/trace"
)

// commentData represents an operation comment formatted to contain tracing data.
type commentData struct {
	DocDB struct {
		TraceID string `json:"traceID"`
		SpanID  string `json:"spanID"`
	} `json:"docdb"`
}

// SpanContextFromComment extracts OpenTelemetry tracing information from an operation comment.
// The comment is expected to be a JSON string (see [commentData]).
//
// If the comment is empty, it returns an invalid span context and no error.
func SpanContextFromComment(comment string) (oteltrace.SpanContext, error) {
	var sc oteltrace.SpanContext

	if comment == "" {
		return sc, nil
	}

	var data commentData

	err := json.Unmarshal([]byte(comment), &data)
	if err != nil {
		return sc, lazyerrors.Error(err)
	}

	traceID, err := oteltrace.TraceIDFromHex(data.DocDB.TraceID)
	if err != nil {
		return sc, lazyerrors.Error(err)
	}

	spanID, err := oteltrace.SpanIDFromHex(data.DocDB.SpanID)
	if err != nil {
		return sc, lazyerrors.Error(err)
	}

	sc = oteltrace.NewSpanContext(oteltrace.SpanContextConfig{
		TraceID:    traceID,
		SpanID:     spanID,
		TraceFlags: oteltrace.FlagsSampled, // we would not have received it if it hadn't been sampled
		Remote:     true,
	})

	return sc, nil
}

// CommentFromSpanContext creates a json-encoded string with tracing information (see commentData) from span context.
func CommentFromSpanContext(sc oteltrace.SpanContext) (string, error) {
	if !sc.IsValid() {
		return "", lazyerrors.New("invalid span context")
	}

	var data commentData
	data.DocDB.TraceID = sc.TraceID().String()
	data.DocDB.SpanID = sc.SpanID().String()

	comment, err := json.Marshal(data)
	if err != nil {
		return "", lazyerrors.Error(err)
	}

	return string(comment), nil
}