Files

338 lines
8.8 KiB
Go

package usage
import (
"bytes"
"context"
"crypto/rand"
"encoding/json"
"io"
"log"
"net/http"
"reflect"
"strconv"
"strings"
"sync"
"sync/atomic"
"time"
"github.com/yetone/magpie/internal/netproxy"
"github.com/yetone/magpie/internal/sessions"
"github.com/yetone/magpie/internal/settings"
)
const otelQueueSize = 128
const otelBatchSize = 32
const otelBodiesBatch = 4
// otelQueueBytes is how much of the queued records' bodies may be held at
// once, so a whole body (#538) can't fill memory: a record past it is
// dropped as the queue's own overflow is.
const otelQueueBytes = 128 << 20
// otelBodiesBatchBytes is how much of one request's bodies go in one OTLP
// POST at most, so a batch of whole bodies stays inside a collector's limit.
const otelBodiesBatchBytes = 16 << 20
var otel atomic.Pointer[otelExporter]
type otelItem struct {
record Record
config settings.OTel
}
type otelExporter struct {
queue chan otelItem
stop, done chan struct{}
ctx context.Context
cancel context.CancelFunc
client *http.Client
salt [8]byte
last time.Time
dropped atomic.Uint64
bytes atomic.Int64 // the queued records' bodies, in bytes
sessions atomic.Pointer[sessionAvailability]
identities sessions.TraceSessionIndex
}
// StartOTel runs only in the process serving the gateway. Stopping drains
// accepted metadata for at most three seconds, after in-flight calls finish.
func StartOTel() func() {
e := newOTelExporter()
otel.Store(e)
if _, err := settings.OTelExport(); err != nil {
log.Printf("otel: %s", err)
}
watchCtx, stopWatch := context.WithCancel(e.ctx)
go e.watchSessions(watchCtx, time.Now())
go e.run()
var once sync.Once
return func() {
once.Do(func() {
stopWatch()
otel.CompareAndSwap(e, nil)
close(e.stop)
timer := time.NewTimer(3 * time.Second)
defer timer.Stop()
select {
case <-e.done:
case <-timer.C:
e.cancel()
<-e.done
}
})
}
}
func newOTelExporter() *otelExporter {
ctx, cancel := context.WithCancel(context.Background())
e := &otelExporter{queue: make(chan otelItem, otelQueueSize), stop: make(chan struct{}), done: make(chan struct{}), ctx: ctx, cancel: cancel,
last: time.Now(), client: &http.Client{Timeout: 3 * time.Second,
Transport: &http.Transport{Proxy: netproxy.Func, MaxIdleConnsPerHost: 2, IdleConnTimeout: 90 * time.Second},
CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }}}
rand.Read(e.salt[:])
return e
}
// ExportOTel queues telemetry without writing a usage ledger record.
func ExportOTel(r Record) { offerOTel(r) }
// OTelEnabled reports whether the running gateway exports traces.
func OTelEnabled() bool {
if otel.Load() == nil {
return false
}
config, err := settings.OTelExport()
return err == nil && config.Enabled
}
func offerOTel(r Record) {
if r.SkipOTel {
return
}
e := otel.Load()
if e == nil {
return
}
config, err := settings.OTelExport()
if err != nil || !config.Enabled {
return
}
session := r.NativeSession
if session == "" {
session = r.Session
}
if r.Local && r.CallerKeyID == "" && r.Via == "" && r.OTel == nil && r.Kind == "" && e.sessionObserved(r.Agent, session, config) {
return
}
if !config.Bodies {
r.BodyIn, r.BodyOut = "", ""
}
e.offer(otelItem{record: r, config: config})
}
// OTelBodies is whether calls' request and reply bodies go with their
// spans (#538): the gateway fills Record.BodyIn/BodyOut only then
func OTelBodies() bool {
if otel.Load() == nil {
return false
}
config, err := settings.OTelExport()
return err == nil && config.Enabled && config.Bodies
}
// OTelWhole is whether calls' bodies go whole, not cut at the gateway's
// 256 KiB Recent-calls capture (#538): the gateway keeps the whole reply for
// the export only then
func OTelWhole() bool {
if otel.Load() == nil {
return false
}
config, err := settings.OTelExport()
return err == nil && config.Enabled && config.Bodies && config.BodiesWhole
}
// recordBytes is how much of a record the queue holds: its bodies, the
// metadata being small beside them
func recordBytes(r Record) int { return len(r.BodyIn) + len(r.BodyOut) }
func (e *otelExporter) offer(item otelItem) {
n := int64(recordBytes(item.record))
if n > 0 {
for {
held := e.bytes.Load()
if n > otelQueueBytes-held {
e.dropped.Add(1) // telemetry must never wait for a slow collector, nor grow without bound
return
}
if e.bytes.CompareAndSwap(held, held+n) {
break
}
}
}
select {
case e.queue <- item:
default:
if n > 0 {
e.bytes.Add(-n)
}
e.dropped.Add(1) // telemetry must never wait for a slow collector
}
}
func (e *otelExporter) run() {
defer close(e.done)
defer e.cancel()
defer e.client.CloseIdleConnections()
ticker := time.NewTicker(5 * time.Second)
defer ticker.Stop()
batch := make([]otelItem, 0, otelBatchSize)
flush := func() {
e.flush(batch)
for _, it := range batch {
if n := int64(recordBytes(it.record)); n > 0 {
e.bytes.Add(-n)
}
}
batch = batch[:0]
}
for {
select {
case it := <-e.queue:
batch = append(batch, it)
if len(batch) == otelBatchSize {
flush()
}
case <-ticker.C:
flush()
case <-e.stop:
for {
select {
case it := <-e.queue:
batch = append(batch, it)
if len(batch) == otelBatchSize {
flush()
}
default:
flush()
return
}
}
}
}
}
func (e *otelExporter) flush(batch []otelItem) {
if len(batch) == 0 {
return
}
config, err := settings.OTelExport()
if err != nil || !config.Enabled {
return
}
var records []Record
for _, it := range batch {
// Never send queued records to a new destination or with changed credentials.
if reflect.DeepEqual(it.config, config) {
records = append(records, it.record)
}
}
if len(records) == 0 {
return
}
now := time.Now()
// spans carrying bodies (up to 256 KiB each way, whole when asked for)
// go a few at a time, so one request stays well inside a collector's
// size limit; a record's own bodies decide when a batch is full too
step := len(records)
if config.Bodies {
step = otelBodiesBatch
}
for i := 0; i < len(records); {
j := min(i+step, len(records))
if config.Bodies {
bytes := 0
for j = i; j < len(records) && j < i+step; j++ {
bytes += recordBytes(records[j])
if j > i && bytes > otelBodiesBatchBytes {
break
}
}
}
e.send(config, "traces", e.traces(records[i:j]))
i = j
}
if config.Metrics {
ledger := make([]Record, 0, len(records))
for _, r := range records {
if r.OTel == nil || !r.OTel.Root && !r.OTel.Update && (!r.OTel.Session || r.OTel.Type == "generation") {
ledger = append(ledger, r)
}
}
if len(ledger) > 0 {
e.send(config, "metrics", otelMetrics(ledger, e.last, now))
}
}
e.last = now
}
func (e *otelExporter) send(config settings.OTel, signal string, payload any) {
b, _ := json.Marshal(payload)
base := strings.TrimRight(config.Endpoint, "/")
base = strings.TrimSuffix(strings.TrimSuffix(base, "/v1/traces"), "/v1/metrics")
endpoint := base + "/v1/" + signal
delay := time.Second
for attempt := 0; attempt < 3; attempt++ {
current, err := settings.OTelExport()
if err != nil || !current.Enabled || !reflect.DeepEqual(current, config) {
return
}
req, err := http.NewRequestWithContext(e.ctx, "POST", endpoint, bytes.NewReader(b))
if err != nil {
return
}
for k, v := range config.Headers {
req.Header.Set(k, v)
}
req.Header.Set("Content-Type", "application/json")
res, err := e.client.Do(req)
retry := err != nil
if res != nil {
response, _ := io.ReadAll(io.LimitReader(res.Body, 64<<10))
res.Body.Close()
if res.StatusCode == http.StatusOK {
var result struct {
PartialSuccess struct {
RejectedSpans json.Number `json:"rejectedSpans"`
RejectedDataPoints json.Number `json:"rejectedDataPoints"`
} `json:"partialSuccess"`
}
json.Unmarshal(response, &result)
if result.PartialSuccess.RejectedSpans != "" && result.PartialSuccess.RejectedSpans != "0" || result.PartialSuccess.RejectedDataPoints != "" && result.PartialSuccess.RejectedDataPoints != "0" {
log.Printf("otel: %s export partially rejected", signal)
}
return
}
retry = res.StatusCode == 429 || res.StatusCode == 502 || res.StatusCode == 503 || res.StatusCode == 504
if n, err := strconv.Atoi(res.Header.Get("Retry-After")); err == nil && n > 0 {
delay = time.Duration(min(n, 60)) * time.Second
} else if at, err := http.ParseTime(res.Header.Get("Retry-After")); err == nil && time.Until(at) > 0 {
delay = time.Until(at)
}
}
if !retry || attempt == 2 {
if e.ctx.Err() == nil {
log.Printf("otel: %s export failed; batch dropped", signal)
}
return
}
timer := time.NewTimer(min(delay, 60*time.Second))
select {
case <-e.ctx.Done():
timer.Stop()
return
case <-timer.C:
}
delay *= 2
}
}