crossbar/internal/admin/admin_ops.go

package admin

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"sort"
	"strconv"
	"strings"
	"time"

	"git.wntrmute.dev/kyle/crossbar/internal/config"
	"git.wntrmute.dev/kyle/crossbar/internal/lease"
	"git.wntrmute.dev/kyle/crossbar/internal/store"
)

// routePin handles POST /_crossbar/routes/{route}: pin the route to a host or
// release and unpin it.
func (hx *handler) routePin(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		wrongMethod(w, "POST")
		return
	}
	route := r.PathValue("route")
	routeCfg, ok := hx.cfg.Routes[route]
	if !ok {
		writeError(w, http.StatusNotFound, "unknown route")
		return
	}

	var raw struct {
		Host    string `json:"host"`
		Pin     *bool  `json:"pin"`
		Release *bool  `json:"release"`
	}
	if err := decodeJSON(r, &raw); err != nil {
		writeError(w, http.StatusBadRequest, "invalid body")
		return
	}

	pinSet := raw.Pin != nil
	releaseSet := raw.Release != nil
	switch {
	case pinSet && releaseSet:
		writeError(w, http.StatusBadRequest, "pin and release at once")
	case !pinSet && !releaseSet:
		writeError(w, http.StatusBadRequest, "pin or release required")
	case pinSet:
		hx.pin(w, route, routeCfg, raw.Host)
	default:
		hx.release(w, route)
	}
}

// pin validates the host, records candidates, and pins the route.
func (hx *handler) pin(w http.ResponseWriter, route string, routeCfg config.Route, host string) {
	if host == "" {
		writeError(w, http.StatusBadRequest, "pin requires host")
		return
	}
	if !containsHost(routeCfg.Hosts, host) {
		writeError(w, http.StatusNotFound, "host not in route")
		return
	}
	// Record the route's hosts as candidates so Pin accepts a host no request
	// has used yet.
	hx.lt.Candidates(route, routeCfg.Hosts)
	if err := hx.lt.Pin(route, host, time.Now()); err != nil {
		if errors.Is(err, lease.ErrUnknownHost) {
			writeError(w, http.StatusNotFound, "unknown host")
			return
		}
		writeError(w, http.StatusInternalServerError, err.Error())
		return
	}
	writeJSON(w, http.StatusOK, map[string]bool{"ok": true})
}

// release drops the route's leases and clears its pin.
func (hx *handler) release(w http.ResponseWriter, route string) {
	n := hx.lt.Release(route)
	hx.lt.Unpin(route)
	writeJSON(w, http.StatusOK, map[string]any{"ok": true, "released": n})
}

// hostDrain handles POST /_crossbar/hosts/{host}: set or clear draining.
func (hx *handler) hostDrain(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		wrongMethod(w, "POST")
		return
	}
	host := r.PathValue("host")
	if _, ok := hx.cfg.Hosts[host]; !ok {
		writeError(w, http.StatusNotFound, "unknown host")
		return
	}
	var body struct {
		Drain *bool `json:"drain"`
	}
	if err := decodeJSON(r, &body); err != nil {
		writeError(w, http.StatusBadRequest, "invalid body")
		return
	}
	if body.Drain == nil {
		writeError(w, http.StatusBadRequest, "drain required")
		return
	}
	hx.d.SetDraining(host, *body.Drain)
	writeJSON(w, http.StatusOK, map[string]bool{"ok": true})
}

// usageGet handles GET /_crossbar/usage: usage rows as JSON, or a fixed-width
// table when Accept is text/plain.
func (hx *handler) usageGet(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		wrongMethod(w, "GET")
		return
	}
	q := r.URL.Query()
	var byv store.By
	switch q.Get("by") {
	case "", "route":
		byv = store.ByRoute
	case "model":
		byv = store.ByModel
	case "host":
		byv = store.ByHost
	default:
		writeError(w, http.StatusBadRequest, "invalid by")
		return
	}
	since, err := parseSince(q.Get("since"))
	if err != nil {
		writeError(w, http.StatusBadRequest, "invalid since")
		return
	}
	rows, err := hx.st.Usage(since, byv)
	if err != nil {
		writeError(w, http.StatusInternalServerError, "usage: "+err.Error())
		return
	}
	if r.Header.Get("Accept") == "text/plain" {
		writeUsageTable(w, rows)
		return
	}
	writeJSON(w, http.StatusOK, rows)
}

// parseSince resolves the since query value: absent means all time, otherwise
// an RFC 3339 instant or a duration (which may end in "d" for days) meaning
// now - d.
func parseSince(s string) (time.Time, error) {
	if s == "" {
		return time.Time{}, nil
	}
	if t, err := time.Parse(time.RFC3339, s); err == nil {
		return t.UTC(), nil
	}
	d, err := parseWindow(s)
	if err != nil {
		return time.Time{}, err
	}
	return time.Now().Add(-d), nil
}

// parseWindow parses a duration, accepting a trailing "d" for whole days.
func parseWindow(s string) (time.Duration, error) {
	if n, ok := splitDays(s); ok {
		return time.Duration(n) * 24 * time.Hour, nil
	}
	return time.ParseDuration(s)
}

// splitDays reports whether s is an integer number of days ("Nd").
func splitDays(s string) (int, bool) {
	if len(s) < 2 || s[len(s)-1] != 'd' {
		return 0, false
	}
	n, err := strconv.Atoi(s[:len(s)-1])
	if err != nil || n < 0 {
		return 0, false
	}
	return n, true
}

// writeUsageTable renders the rows as a fixed-width table with a header line,
// one row per entry, no trailing spaces.
func writeUsageTable(w http.ResponseWriter, rows []store.UsageRow) {
	headers := []string{"key", "requests", "errors", "busy_ms", "queued_ms", "prompt", "cached", "completion", "cache_hit"}
	lines := make([][]string, 0, len(rows)+1)
	lines = append(lines, headers)
	for _, u := range rows {
		lines = append(lines, []string{
			u.Key,
			strconv.FormatInt(u.Requests, 10),
			strconv.FormatInt(u.Errors, 10),
			strconv.FormatInt(u.BusyMs, 10),
			strconv.FormatInt(u.QueuedMs, 10),
			strconv.FormatInt(u.PromptTokens, 10),
			strconv.FormatInt(u.CachedTokens, 10),
			strconv.FormatInt(u.CompletionTokens, 10),
			strconv.FormatFloat(u.CacheHitRatio(), 'f', 2, 64),
		})
	}
	widths := columnWidths(lines)

	var b strings.Builder
	for _, line := range lines {
		for i, f := range line {
			if i < len(line)-1 {
				b.WriteString(fmt.Sprintf("%-*s ", widths[i], f))
			} else {
				b.WriteString(f)
			}
		}
		b.WriteByte('\n')
	}
	w.Header().Set("Content-Type", "text/plain; charset=utf-8")
	w.WriteHeader(http.StatusOK)
	_, _ = w.Write([]byte(b.String()))
}

// columnWidths returns the widest rendered field in each column.
func columnWidths(lines [][]string) []int {
	widths := make([]int, len(lines[0]))
	for _, line := range lines {
		for i, f := range line {
			if len(f) > widths[i] {
				widths[i] = len(f)
			}
		}
	}
	return widths
}

// metricsGet handles GET /_crossbar/metrics, emitting the Prometheus text
// exposition format computed on request.
func (hx *handler) metricsGet(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		wrongMethod(w, "GET")
		return
	}
	counts, err := hx.st.StatusCounts(time.Time{})
	if err != nil {
		writeError(w, http.StatusInternalServerError, "metrics: "+err.Error())
		return
	}
	usage, err := hx.st.Usage(time.Time{}, store.ByRoute)
	if err != nil {
		writeError(w, http.StatusInternalServerError, "metrics: "+err.Error())
		return
	}

	var reqSamples []string
	for _, c := range counts {
		reqSamples = append(reqSamples, fmt.Sprintf(
			"crossbar_requests_total{route=\"%s\",host=\"%s\",status=\"%s\"} %d",
			esc(c.Route), esc(c.Host), esc(strconv.Itoa(c.Status)), c.Count))
	}

	var prompt, cached, queue []string
	for _, u := range usage {
		prompt = append(prompt, fmt.Sprintf("crossbar_prompt_tokens_total{route=\"%s\"} %d", esc(u.Key), u.PromptTokens))
		cached = append(cached, fmt.Sprintf("crossbar_cached_tokens_total{route=\"%s\"} %d", esc(u.Key), u.CachedTokens))
		queue = append(queue, fmt.Sprintf("crossbar_queue_wait_ms_total{route=\"%s\"} %d", esc(u.Key), u.QueuedMs))
	}

	all := hx.h.All()
	names := make([]string, 0, len(all))
	for name := range all {
		names = append(names, name)
	}
	sort.Strings(names)

	var healthy, free, inflight, queued []string
	for _, name := range names {
		s := all[name]
		healthy = append(healthy, fmt.Sprintf("crossbar_host_healthy{host=\"%s\"} %d", esc(name), btoi(s.Healthy)))
		free = append(free, fmt.Sprintf("crossbar_host_free_slots{host=\"%s\"} %d", esc(name), hx.lim.FreeSlots(name)))
		fi, q := 0, 0
		for _, m := range configuredModels(hx.cfg, name) {
			fi += hx.lim.InFlight(name, m)
			q += hx.lim.Queued(name, m)
		}
		inflight = append(inflight, fmt.Sprintf("crossbar_host_in_flight{host=\"%s\"} %d", esc(name), fi))
		queued = append(queued, fmt.Sprintf("crossbar_host_queued{host=\"%s\"} %d", esc(name), q))
	}

	var b strings.Builder
	appendFamily(&b, "crossbar_requests_total", "counter", reqSamples)
	appendFamily(&b, "crossbar_prompt_tokens_total", "counter", prompt)
	appendFamily(&b, "crossbar_cached_tokens_total", "counter", cached)
	appendFamily(&b, "crossbar_queue_wait_ms_total", "counter", queue)
	appendFamily(&b, "crossbar_host_healthy", "gauge", healthy)
	appendFamily(&b, "crossbar_host_free_slots", "gauge", free)
	appendFamily(&b, "crossbar_host_in_flight", "gauge", inflight)
	appendFamily(&b, "crossbar_host_queued", "gauge", queued)

	w.Header().Set("Content-Type", "text/plain; version=0.0.4")
	w.WriteHeader(http.StatusOK)
	_, _ = w.Write([]byte(b.String()))
}

// appendFamily writes a metric family: its TYPE line followed by the sorted
// sample lines.
func appendFamily(b *strings.Builder, name, typ string, samples []string) {
	fmt.Fprintf(b, "# TYPE %s %s\n", name, typ)
	sort.Strings(samples)
	for _, s := range samples {
		b.WriteString(s)
		b.WriteByte('\n')
	}
}

// esc escapes a label value for the Prometheus text format.
func esc(s string) string {
	s = strings.ReplaceAll(s, `\`, `\\`)
	s = strings.ReplaceAll(s, `"`, `\"`)
	return s
}

// btoi converts a bool to 0/1 for a gauge.
func btoi(v bool) int {
	if v {
		return 1
	}
	return 0
}

// containsHost reports whether hosts contains h.
func containsHost(hosts []string, h string) bool {
	for _, x := range hosts {
		if x == h {
			return true
		}
	}
	return false
}

// decodeJSON decodes a bounded JSON body.
func decodeJSON(r *http.Request, v any) error {
	dec := json.NewDecoder(io.LimitReader(r.Body, 4096))
	return dec.Decode(v)
}

func writeJSON(w http.ResponseWriter, status int, v any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	_ = json.NewEncoder(w).Encode(v)
}

func writeError(w http.ResponseWriter, status int, msg string) {
	writeJSON(w, status, map[string]string{"error": msg})
}

// wrongMethod answers 405 with the allowed method in the Allow header.
func wrongMethod(w http.ResponseWriter, allow string) {
	w.Header().Set("Allow", allow)
	writeJSON(w, http.StatusMethodNotAllowed, map[string]string{"error": "method not allowed"})
}

func (hx *handler) unknown(w http.ResponseWriter, r *http.Request) {
	writeJSON(w, http.StatusNotFound, map[string]string{"error": "not found"})
}