orbit/internal/observability/http.go

package observability

import (
	"bufio"
	"io"
	"net"
	"net/http"
	"strconv"
	"strings"
	"time"
)

// Middleware counts HTTP completion once. A successful upgrade completes HTTP
// at Hijack; the still-running WebSocket handler is counted separately.
func (r *Registry) Middleware(next http.Handler) http.Handler {
	if r == nil {
		return next
	}
	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		started := time.Now()
		route := routeClass(req.URL.Path)
		r.AddGauge("orbit_http_requests_in_flight", 1, nil)
		cw := &captureWriter{ResponseWriter: w}
		completed := false
		ws := false
		complete := func(status int) {
			if completed {
				return
			}
			completed = true
			class := "other"
			if status >= 100 && status < 600 {
				class = strconv.Itoa(status/100) + "xx"
			}
			labels := Labels{"method": boundedMethod(req.Method), "route": route, "status_class": class}
			r.AddGauge("orbit_http_requests_in_flight", -1, nil)
			r.AddCounter("orbit_http_requests_total", 1, labels)
			r.ObserveHistogram("orbit_http_request_duration_seconds", time.Since(started).Seconds(), labels)
		}
		cw.onHijack = func() {
			complete(cw.status)
			if cw.status == 101 && (route == "worker_connect" || route == "event_connect") {
				ws = true
				r.AddGauge("orbit_websocket_connections", 1, Labels{"route": route})
				r.AddCounter("orbit_websocket_connections_total", 1, Labels{"route": route})
			}
		}
		returned := false
		defer func() {
			status := cw.status
			if !returned {
				r.AddCounter("orbit_http_panics_total", 1, nil)
				if status == 0 {
					status = 500
				}
			}
			if status == 0 {
				status = 200
			}
			complete(status)
			if ws {
				r.AddGauge("orbit_websocket_connections", -1, Labels{"route": route})
			}
			// No recover: net/http retains its panic/ErrAbortHandler behavior.
		}()
		next.ServeHTTP(cw, req)
		returned = true
	})
}
func boundedMethod(method string) string {
	for _, v := range methods {
		if method == v {
			return v
		}
	}
	return "OTHER"
}
func routeClass(path string) string {
	switch path {
	case "/health/live":
		return "health_live"
	case "/health/ready":
		return "health_ready"
	case "/v1/worker/connect":
		return "worker_connect"
	case "/v1/events/connect":
		return "event_connect"
	}
	for _, p := range []struct{ prefix, route string }{
		{"/v1/tasks", "tasks"}, {"/v1/queues", "queues"}, {"/v1/workflows", "workflows"},
		{"/v1/workflow-definitions", "workflow_definitions"}, {"/v1/schedules", "schedules"},
		{"/v1/subscriptions", "subscriptions"}, {"/v1/application/resources", "application_resources"},
	} {
		if path == p.prefix || strings.HasPrefix(path, p.prefix+"/") || strings.HasPrefix(path, p.prefix+":") {
			return p.route
		}
	}
	return "other"
}

type captureWriter struct {
	http.ResponseWriter
	status   int
	onHijack func()
}

func (w *captureWriter) WriteHeader(status int) {
	if w.status != 0 {
		return
	}
	// Forward informational responses without preventing a later final response.
	if status >= 100 && status < 200 && status != 101 {
		w.ResponseWriter.WriteHeader(status)
		return
	}
	w.ResponseWriter.WriteHeader(status)
	w.status = status
}
func (w *captureWriter) Write(p []byte) (int, error) {
	if w.status == 0 {
		w.WriteHeader(200)
	}
	return w.ResponseWriter.Write(p)
}
func (w *captureWriter) Unwrap() http.ResponseWriter { return w.ResponseWriter }
func (w *captureWriter) FlushError() error {
	if w.status == 0 {
		w.WriteHeader(200)
	}
	return http.NewResponseController(w.ResponseWriter).Flush()
}
func (w *captureWriter) Flush() { _ = w.FlushError() }
func (w *captureWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	c, b, err := http.NewResponseController(w.ResponseWriter).Hijack()
	if err == nil && w.onHijack != nil {
		w.onHijack()
	}
	return c, b, err
}
func (w *captureWriter) Push(target string, opts *http.PushOptions) error {
	if p, ok := w.ResponseWriter.(http.Pusher); ok {
		return p.Push(target, opts)
	}
	return http.ErrNotSupported
}
func (w *captureWriter) ReadFrom(src io.Reader) (int64, error) {
	if w.status == 0 {
		w.WriteHeader(200)
	}
	if rf, ok := w.ResponseWriter.(io.ReaderFrom); ok {
		return rf.ReadFrom(src)
	}
	return io.Copy(w.ResponseWriter, src)
}