llmgw/internal/control/reconciler/trace.go

package reconciler

import (
	"context"
	"errors"
	"sync"
	"time"

	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/codes"
	"go.opentelemetry.io/otel/trace"

	"gitea.mynodes.xyz/ilker/llmgw/internal/obs"
)

// What is traced here, and why it is shaped this way.
//
// A pass runs every few seconds and nearly always finds nothing to do. A
// trace per pass would bury the few that placed or stopped something under
// seventeen thousand empty ones a day, so a pass's span is begun only when
// the pass first does something, and backdated to when the pass began.
//
// A replica's start outlives the pass that placed it by minutes. Its span,
// "replica start", is opened beneath the placement and closed when the
// replica serves, fails, or is stopped, and the command to each node carries
// it, so the agent's spans for that shard — weights, launch, readiness — are
// its descendants. That one trace is the answer to "why did this model take
// four minutes to start".

const tracerName = "gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"

func tracer() trace.Tracer { return obs.Tracer(tracerName) }

// pass is one reconcile pass's span, begun on first use.
type pass struct {
	start time.Time
	base  context.Context

	once sync.Once
	ctx  context.Context
	span trace.Span
}

type passKey struct{}

// beginPass marks ctx as a pass's, for startSpan to hang its spans on.
func beginPass(ctx context.Context) (context.Context, *pass) {
	p := &pass{start: time.Now(), base: ctx}
	return context.WithValue(ctx, passKey{}, p), p
}

// context is the pass's span's context, beginning the span if it has not
// been.
func (p *pass) context() context.Context {
	p.once.Do(func() {
		p.ctx, p.span = tracer().Start(p.base, "reconcile", trace.WithTimestamp(p.start))
	})
	return p.ctx
}

// end closes the pass's span, if anything began it.
func (p *pass) end() {
	p.once.Do(func() {}) // nothing may begin it after this
	if p.span != nil {
		p.span.End()
	}
}

// startSpan opens a span for something the reconciler did: beneath the span
// already in ctx, or beneath the pass's — begun now — when ctx holds the
// pass and nothing else.
func startSpan(ctx context.Context, name string, attrs ...attribute.KeyValue) (context.Context, trace.Span) {
	if !trace.SpanContextFromContext(ctx).IsValid() {
		if p, ok := ctx.Value(passKey{}).(*pass); ok {
			if pctx := p.context(); pctx != nil {
				ctx = pctx
			}
		}
	}
	return tracer().Start(ctx, name, trace.WithAttributes(attrs...))
}

// endSpan closes a span, marking it failed when why is not empty.
func endSpan(span trace.Span, why string) {
	if why != "" {
		obs.End(span, errors.New(why))
		return
	}
	span.End()
}

// replicaAttrs are what every span about one replica says.
func replicaAttrs(rep *Replica) []attribute.KeyValue {
	return []attribute.KeyValue{
		attribute.String("llmgw.replica", rep.ID),
		attribute.String("llmgw.model", rep.Model),
		attribute.String("llmgw.profile", rep.Profile),
		attribute.StringSlice("llmgw.nodes", rep.NodeIDs()),
	}
}

// startDone closes a replica's "replica start" span, once: when it serves,
// with how long that took, or when it fails or is stopped first, with why.
// Called with r.mu held, or before the replica is in the map.
func (rep *Replica) startDone(now time.Time) {
	if rep.startSpan == nil {
		return
	}
	span := rep.startSpan
	rep.startSpan = nil
	span.SetAttributes(attribute.String("llmgw.state", rep.State.String()))
	switch rep.State {
	case Ready:
		span.SetAttributes(attribute.Float64("llmgw.ready_after_seconds", now.Sub(rep.CreatedAt).Seconds()))
		span.SetStatus(codes.Ok, "")
	case Failed:
		span.SetStatus(codes.Error, rep.Err)
	default:
		span.SetStatus(codes.Error, "stopped before it was ready")
	}
	span.End()
}