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) }