-
This commit is contained in:
@@ -0,0 +1,376 @@
|
||||
package telemetry
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/example/ollama-fair-gateway/internal/config"
|
||||
)
|
||||
|
||||
type Context struct {
|
||||
TraceID string
|
||||
RootSpanID string
|
||||
ParentSpanID string
|
||||
Sampled bool
|
||||
}
|
||||
|
||||
type Interval struct {
|
||||
Name string
|
||||
Start time.Time
|
||||
End time.Time
|
||||
Attrs map[string]any
|
||||
}
|
||||
|
||||
type Record struct {
|
||||
Trace Context
|
||||
RequestID string
|
||||
API string
|
||||
Path string
|
||||
Tenant string
|
||||
Actor string
|
||||
Application string
|
||||
ServiceClass string
|
||||
Model string
|
||||
Alias string
|
||||
Worker string
|
||||
Started time.Time
|
||||
Finished time.Time
|
||||
FirstByte time.Duration
|
||||
Status int
|
||||
PromptTokens int64
|
||||
OutputTokens int64
|
||||
CachedTokens int64
|
||||
Credits float64
|
||||
Error string
|
||||
Intervals []Interval
|
||||
}
|
||||
|
||||
type Exporter struct {
|
||||
cfg config.OpenTelemetryConfig
|
||||
client *http.Client
|
||||
ch chan []span
|
||||
stop chan struct{}
|
||||
done chan struct{}
|
||||
closed atomic.Bool
|
||||
dropped atomic.Uint64
|
||||
failed atomic.Uint64
|
||||
exported atomic.Uint64
|
||||
}
|
||||
|
||||
type span struct {
|
||||
TraceID string
|
||||
SpanID string
|
||||
ParentSpanID string
|
||||
Name string
|
||||
Kind int
|
||||
Start time.Time
|
||||
End time.Time
|
||||
Attrs map[string]any
|
||||
Error string
|
||||
}
|
||||
|
||||
func New(cfg config.OpenTelemetryConfig) *Exporter {
|
||||
if !cfg.Enabled {
|
||||
return nil
|
||||
}
|
||||
if cfg.BatchSize <= 0 {
|
||||
cfg.BatchSize = 128
|
||||
}
|
||||
if cfg.FlushInterval == 0 {
|
||||
cfg.FlushInterval = config.Duration(2 * time.Second)
|
||||
}
|
||||
e := &Exporter{cfg: cfg, client: &http.Client{Timeout: 10 * time.Second}, ch: make(chan []span, cfg.BatchSize*4), stop: make(chan struct{}), done: make(chan struct{})}
|
||||
go e.loop()
|
||||
return e
|
||||
}
|
||||
|
||||
func (e *Exporter) NewTrace(traceparent string) Context {
|
||||
if e == nil {
|
||||
return Context{}
|
||||
}
|
||||
ratio := e.cfg.SampleRatio
|
||||
if ratio <= 0 || ratio > 1 {
|
||||
ratio = 1
|
||||
}
|
||||
sampled := ratio >= 1 || randomFraction() < ratio
|
||||
traceID, parent := parseTraceParent(traceparent)
|
||||
if traceID == "" {
|
||||
traceID = randomHex(16)
|
||||
}
|
||||
return Context{TraceID: traceID, RootSpanID: randomHex(8), ParentSpanID: parent, Sampled: sampled}
|
||||
}
|
||||
func (e *Exporter) TraceParent(c Context) string {
|
||||
if e == nil || !c.Sampled {
|
||||
return ""
|
||||
}
|
||||
return "00-" + c.TraceID + "-" + c.RootSpanID + "-01"
|
||||
}
|
||||
func (e *Exporter) Dropped() uint64 {
|
||||
if e == nil {
|
||||
return 0
|
||||
}
|
||||
return e.dropped.Load()
|
||||
}
|
||||
func (e *Exporter) Failed() uint64 {
|
||||
if e == nil {
|
||||
return 0
|
||||
}
|
||||
return e.failed.Load()
|
||||
}
|
||||
func (e *Exporter) Exported() uint64 {
|
||||
if e == nil {
|
||||
return 0
|
||||
}
|
||||
return e.exported.Load()
|
||||
}
|
||||
|
||||
func (e *Exporter) Record(r Record) {
|
||||
if e == nil || !r.Trace.Sampled || e.closed.Load() {
|
||||
return
|
||||
}
|
||||
if r.Finished.IsZero() {
|
||||
r.Finished = time.Now().UTC()
|
||||
}
|
||||
attrs := map[string]any{
|
||||
"gen_ai.operation.name": operationName(r.Path),
|
||||
"gen_ai.provider.name": "ollama",
|
||||
"gen_ai.request.model": r.Model,
|
||||
"gen_ai.usage.input_tokens": r.PromptTokens,
|
||||
"gen_ai.usage.output_tokens": r.OutputTokens,
|
||||
"gen_ai.usage.cache_read.input_tokens": r.CachedTokens,
|
||||
"http.request.method": "POST",
|
||||
"http.response.status_code": r.Status,
|
||||
"server.address": r.Worker,
|
||||
"ollama.gateway.request_id": r.RequestID,
|
||||
"ollama.gateway.api": r.API,
|
||||
"ollama.gateway.tenant": r.Tenant,
|
||||
"ollama.gateway.actor": r.Actor,
|
||||
"ollama.gateway.application": r.Application,
|
||||
"ollama.gateway.service_class": r.ServiceClass,
|
||||
"ollama.gateway.credits": r.Credits,
|
||||
}
|
||||
if r.Alias != "" {
|
||||
attrs["ollama.gateway.model_alias"] = r.Alias
|
||||
}
|
||||
if r.FirstByte > 0 {
|
||||
attrs["gen_ai.response.time_to_first_chunk"] = r.FirstByte.Seconds()
|
||||
}
|
||||
root := span{TraceID: r.Trace.TraceID, SpanID: r.Trace.RootSpanID, ParentSpanID: r.Trace.ParentSpanID, Name: operationName(r.Path) + " " + r.Model, Kind: 2, Start: r.Started, End: r.Finished, Attrs: attrs, Error: r.Error}
|
||||
spans := []span{root}
|
||||
for _, iv := range r.Intervals {
|
||||
if iv.Start.IsZero() || iv.End.IsZero() || iv.End.Before(iv.Start) {
|
||||
continue
|
||||
}
|
||||
a := map[string]any{"ollama.gateway.request_id": r.RequestID}
|
||||
for k, v := range iv.Attrs {
|
||||
a[k] = v
|
||||
}
|
||||
spans = append(spans, span{TraceID: r.Trace.TraceID, SpanID: randomHex(8), ParentSpanID: r.Trace.RootSpanID, Name: iv.Name, Kind: 1, Start: iv.Start, End: iv.End, Attrs: a})
|
||||
}
|
||||
select {
|
||||
case e.ch <- spans:
|
||||
default:
|
||||
e.dropped.Add(uint64(len(spans)))
|
||||
}
|
||||
}
|
||||
|
||||
func (e *Exporter) Close(ctx context.Context) error {
|
||||
if e == nil || !e.closed.CompareAndSwap(false, true) {
|
||||
return nil
|
||||
}
|
||||
close(e.stop)
|
||||
select {
|
||||
case <-e.done:
|
||||
return nil
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
func (e *Exporter) loop() {
|
||||
defer close(e.done)
|
||||
tick := time.NewTicker(e.cfg.FlushInterval.Value())
|
||||
defer tick.Stop()
|
||||
batch := make([]span, 0, e.cfg.BatchSize)
|
||||
flush := func() {
|
||||
if len(batch) == 0 {
|
||||
return
|
||||
}
|
||||
if err := e.export(batch); err != nil {
|
||||
e.failed.Add(uint64(len(batch)))
|
||||
} else {
|
||||
e.exported.Add(uint64(len(batch)))
|
||||
}
|
||||
batch = batch[:0]
|
||||
}
|
||||
for {
|
||||
select {
|
||||
case xs := <-e.ch:
|
||||
batch = append(batch, xs...)
|
||||
if len(batch) >= e.cfg.BatchSize {
|
||||
flush()
|
||||
}
|
||||
case <-tick.C:
|
||||
flush()
|
||||
case <-e.stop:
|
||||
for {
|
||||
select {
|
||||
case xs := <-e.ch:
|
||||
batch = append(batch, xs...)
|
||||
default:
|
||||
flush()
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (e *Exporter) export(spans []span) error {
|
||||
endpoint := strings.TrimSpace(e.cfg.Endpoint)
|
||||
if endpoint == "" {
|
||||
return fmt.Errorf("empty OTLP endpoint")
|
||||
}
|
||||
u, err := url.Parse(endpoint)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if u.Path == "" || u.Path == "/" {
|
||||
u.Path = "/v1/traces"
|
||||
} else if !strings.HasSuffix(u.Path, "/v1/traces") {
|
||||
u.Path = strings.TrimRight(u.Path, "/") + "/v1/traces"
|
||||
}
|
||||
payload := map[string]any{"resourceSpans": []any{map[string]any{"resource": map[string]any{"attributes": attrsToOTLP(map[string]any{"service.name": e.cfg.ServiceName, "service.version": e.cfg.ServiceVersion, "telemetry.sdk.name": "ollama-fair-gateway", "telemetry.sdk.language": "go"})}, "scopeSpans": []any{map[string]any{"scope": map[string]any{"name": "ollama-fair-gateway"}, "spans": spansToOTLP(spans)}}}}}
|
||||
b, _ := json.Marshal(payload)
|
||||
req, err := http.NewRequest(http.MethodPost, u.String(), bytes.NewReader(b))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
for k, v := range e.cfg.Headers {
|
||||
req.Header.Set(k, v)
|
||||
}
|
||||
resp, err := e.client.Do(req)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
_, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 4<<20))
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
return fmt.Errorf("OTLP HTTP %d", resp.StatusCode)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func spansToOTLP(in []span) []any {
|
||||
out := make([]any, 0, len(in))
|
||||
for _, s := range in {
|
||||
x := map[string]any{"traceId": s.TraceID, "spanId": s.SpanID, "name": s.Name, "kind": s.Kind, "startTimeUnixNano": fmt.Sprintf("%d", s.Start.UnixNano()), "endTimeUnixNano": fmt.Sprintf("%d", s.End.UnixNano()), "attributes": attrsToOTLP(s.Attrs)}
|
||||
if s.ParentSpanID != "" {
|
||||
x["parentSpanId"] = s.ParentSpanID
|
||||
}
|
||||
if s.Error != "" {
|
||||
x["status"] = map[string]any{"code": 2, "message": s.Error}
|
||||
}
|
||||
out = append(out, x)
|
||||
}
|
||||
return out
|
||||
}
|
||||
func attrsToOTLP(in map[string]any) []any {
|
||||
keys := make([]string, 0, len(in))
|
||||
for k, v := range in {
|
||||
if !empty(v) {
|
||||
keys = append(keys, k)
|
||||
}
|
||||
}
|
||||
sortStrings(keys)
|
||||
out := make([]any, 0, len(keys))
|
||||
for _, k := range keys {
|
||||
out = append(out, map[string]any{"key": k, "value": otlpValue(in[k])})
|
||||
}
|
||||
return out
|
||||
}
|
||||
func otlpValue(v any) map[string]any {
|
||||
switch x := v.(type) {
|
||||
case string:
|
||||
return map[string]any{"stringValue": x}
|
||||
case bool:
|
||||
return map[string]any{"boolValue": x}
|
||||
case int:
|
||||
return map[string]any{"intValue": fmt.Sprintf("%d", x)}
|
||||
case int64:
|
||||
return map[string]any{"intValue": fmt.Sprintf("%d", x)}
|
||||
case uint64:
|
||||
return map[string]any{"intValue": fmt.Sprintf("%d", x)}
|
||||
case float64:
|
||||
return map[string]any{"doubleValue": x}
|
||||
default:
|
||||
return map[string]any{"stringValue": fmt.Sprint(x)}
|
||||
}
|
||||
}
|
||||
func empty(v any) bool {
|
||||
switch x := v.(type) {
|
||||
case string:
|
||||
return x == ""
|
||||
case int:
|
||||
return x == 0
|
||||
case int64:
|
||||
return x == 0
|
||||
case uint64:
|
||||
return x == 0
|
||||
case float64:
|
||||
return x == 0
|
||||
}
|
||||
return false
|
||||
}
|
||||
func operationName(path string) string {
|
||||
p := strings.ToLower(path)
|
||||
switch {
|
||||
case strings.Contains(p, "embed"):
|
||||
return "embeddings"
|
||||
case strings.Contains(p, "chat") || strings.Contains(p, "messages"):
|
||||
return "chat"
|
||||
default:
|
||||
return "text_completion"
|
||||
}
|
||||
}
|
||||
func randomHex(n int) string { b := make([]byte, n); _, _ = rand.Read(b); return hex.EncodeToString(b) }
|
||||
func randomFraction() float64 {
|
||||
b := make([]byte, 8)
|
||||
_, _ = rand.Read(b)
|
||||
var x uint64
|
||||
for _, z := range b {
|
||||
x = x<<8 | uint64(z)
|
||||
}
|
||||
return float64(x>>11) / float64(uint64(1)<<53)
|
||||
}
|
||||
func parseTraceParent(v string) (string, string) {
|
||||
parts := strings.Split(strings.TrimSpace(v), "-")
|
||||
if len(parts) != 4 || len(parts[1]) != 32 || len(parts[2]) != 16 {
|
||||
return "", ""
|
||||
}
|
||||
if _, e := hex.DecodeString(parts[1]); e != nil {
|
||||
return "", ""
|
||||
}
|
||||
if _, e := hex.DecodeString(parts[2]); e != nil {
|
||||
return "", ""
|
||||
}
|
||||
return strings.ToLower(parts[1]), strings.ToLower(parts[2])
|
||||
}
|
||||
func sortStrings(x []string) {
|
||||
for i := 1; i < len(x); i++ {
|
||||
for j := i; j > 0 && x[j] < x[j-1]; j-- {
|
||||
x[j], x[j-1] = x[j-1], x[j]
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,66 @@
|
||||
package telemetry
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/example/ollama-fair-gateway/internal/config"
|
||||
)
|
||||
|
||||
func TestOTLPHTTPJSONExportIsMetadataOnly(t *testing.T) {
|
||||
bodyCh := make(chan []byte, 1)
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/v1/traces" {
|
||||
t.Fatalf("path=%s", r.URL.Path)
|
||||
}
|
||||
if r.Header.Get("Content-Type") != "application/json" {
|
||||
t.Fatalf("content-type=%s", r.Header.Get("Content-Type"))
|
||||
}
|
||||
b, _ := io.ReadAll(r.Body)
|
||||
bodyCh <- b
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
e := New(config.OpenTelemetryConfig{Enabled: true, Endpoint: ts.URL, ServiceName: "gateway-test", SampleRatio: 1, BatchSize: 1, FlushInterval: config.Duration(10 * time.Millisecond)})
|
||||
ctx := e.NewTrace("00-0123456789abcdef0123456789abcdef-0123456789abcdef-01")
|
||||
start := time.Now().Add(-time.Second)
|
||||
e.Record(Record{Trace: ctx, RequestID: "req-1", API: "openai", Path: "/v1/chat/completions", Tenant: "t", Actor: "u", ServiceClass: "interactive", Model: "qwen3:8b", Worker: "gpu", Started: start, Finished: time.Now(), FirstByte: 100 * time.Millisecond, Status: 200, PromptTokens: 12, OutputTokens: 7, Credits: 2.5, Intervals: []Interval{{Name: "gateway.queue", Start: start, End: start.Add(20 * time.Millisecond)}}})
|
||||
|
||||
var b []byte
|
||||
select {
|
||||
case b = <-bodyCh:
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("no OTLP export")
|
||||
}
|
||||
if strings.Contains(string(b), "secret prompt") {
|
||||
t.Fatal("content leaked")
|
||||
}
|
||||
var doc map[string]any
|
||||
if err := json.Unmarshal(b, &doc); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(doc["resourceSpans"].([]any)) != 1 {
|
||||
t.Fatalf("bad resourceSpans: %s", b)
|
||||
}
|
||||
if !strings.Contains(string(b), "gen_ai.usage.input_tokens") || !strings.Contains(string(b), "gen_ai.response.time_to_first_chunk") {
|
||||
t.Fatalf("missing GenAI attrs: %s", b)
|
||||
}
|
||||
if !strings.Contains(string(b), "0123456789abcdef0123456789abcdef") {
|
||||
t.Fatalf("incoming trace id not propagated: %s", b)
|
||||
}
|
||||
cctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
if err := e.Close(cctx); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if e.Exported() == 0 {
|
||||
t.Fatal("exported counter not incremented")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user