51 lines
1.2 KiB
Go
51 lines
1.2 KiB
Go
|
package server
|
||
|
|
|
||
|
|
import (
|
||
|
|
"log/slog"
|
||
|
|
"net/http"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
"g.fc.run/ai/byre/internal/trace"
|
||
|
|
)
|
||
|
|
|
||
|
|
const slowRequest = 200 * time.Millisecond
|
||
|
|
|
||
|
|
type timedWriter struct {
|
||
|
|
http.ResponseWriter
|
||
|
|
start time.Time
|
||
|
|
trace *trace.Trace
|
||
|
|
status int
|
||
|
|
}
|
||
|
|
|
||
|
|
func (w *timedWriter) WriteHeader(status int) {
|
||
|
|
if w.status == 0 {
|
||
|
|
w.status = status
|
||
|
|
w.Header().Set("Server-Timing", w.trace.Header(time.Since(w.start)))
|
||
|
|
}
|
||
|
|
w.ResponseWriter.WriteHeader(status)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (w *timedWriter) Write(b []byte) (int, error) {
|
||
|
|
if w.status == 0 {
|
||
|
|
w.WriteHeader(http.StatusOK)
|
||
|
|
}
|
||
|
|
return w.ResponseWriter.Write(b)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (w *timedWriter) Unwrap() http.ResponseWriter {
|
||
|
|
return w.ResponseWriter
|
||
|
|
}
|
||
|
|
|
||
|
|
func timing(next http.Handler) http.Handler {
|
||
|
|
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
|
|
ctx, t := trace.New(r.Context())
|
||
|
|
tw := &timedWriter{ResponseWriter: w, start: time.Now(), trace: t}
|
||
|
|
next.ServeHTTP(tw, r.WithContext(ctx))
|
||
|
|
took := time.Since(tw.start)
|
||
|
|
if took < slowRequest || tw.status == http.StatusSwitchingProtocols || tw.Header().Get("Content-Type") == "text/event-stream" {
|
||
|
|
return
|
||
|
|
}
|
||
|
|
slog.Info("request slow", "method", r.Method, "path", r.URL.Path, "status", tw.status, "took", took, "spans", t.String())
|
||
|
|
})
|
||
|
|
}
|