im_server/server/monitoring_http.go

package main

import (
	"bufio"
	"context"
	"encoding/json"
	"io"
	"net"
	"net/http"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/prometheus/client_golang/prometheus/promhttp"
	immetrics "github.com/tinode/chat/server/metrics"
	"github.com/tinode/chat/server/store"
)

var (
	metricsHandler        = promhttp.HandlerFor(immetrics.Registry, promhttp.HandlerOpts{})
	uuidLikeMetricSegment = regexp.MustCompile(`^[0-9a-fA-F]{8}-[0-9a-fA-F-]{13,}$`)
)

type healthzResponse struct {
	Status string            `json:"status"`
	Checks map[string]string `json:"checks"`
}

func serveHealthz(wrt http.ResponseWriter, req *http.Request) {
	if req.Method != http.MethodGet && req.Method != http.MethodHead {
		wrt.Header().Set("Allow", http.MethodGet+", "+http.MethodHead)
		http.Error(wrt, "method not allowed", http.StatusMethodNotAllowed)
		return
	}
	wrt.Header().Set("Content-Type", "application/json; charset=utf-8")

	checks := map[string]string{}
	status := http.StatusOK
	ctx, cancel := context.WithTimeout(req.Context(), 2*time.Second)
	defer cancel()
	if err := store.Store.HealthCheck(ctx); err != nil {
		checks["mysql"] = err.Error()
		status = http.StatusServiceUnavailable
	} else {
		checks["mysql"] = "ok"
	}

	body := healthzResponse{Status: "ok", Checks: checks}
	if status != http.StatusOK {
		body.Status = "error"
	}
	wrt.WriteHeader(status)
	if req.Method != http.MethodHead {
		_ = json.NewEncoder(wrt).Encode(body)
	}
}

func serveMetrics(wrt http.ResponseWriter, req *http.Request) {
	if req.Method != http.MethodGet && req.Method != http.MethodHead {
		wrt.Header().Set("Allow", http.MethodGet+", "+http.MethodHead)
		http.Error(wrt, "method not allowed", http.StatusMethodNotAllowed)
		return
	}
	metricsHandler.ServeHTTP(wrt, req)
}

func monitorHTTP(next http.Handler) http.Handler {
	return http.HandlerFunc(func(wrt http.ResponseWriter, req *http.Request) {
		if req.URL.Path == "/metrics" {
			next.ServeHTTP(wrt, req)
			return
		}
		started := time.Now()
		rec := &monitoringResponseWriter{ResponseWriter: wrt, status: http.StatusOK}
		next.ServeHTTP(rec, req)
		status := strconv.Itoa(rec.status)
		path := normalizeMetricPath(req.URL.Path)
		immetrics.HTTPRequestsTotal.WithLabelValues(req.Method, path, status).Inc()
		immetrics.HTTPRequestDuration.WithLabelValues(req.Method, path, status).Observe(time.Since(started).Seconds())
	})
}

type monitoringResponseWriter struct {
	http.ResponseWriter
	status int
}

func (w *monitoringResponseWriter) WriteHeader(status int) {
	w.status = status
	w.ResponseWriter.WriteHeader(status)
}

func (w *monitoringResponseWriter) Flush() {
	if f, ok := w.ResponseWriter.(http.Flusher); ok {
		f.Flush()
	}
}

func (w *monitoringResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	h, ok := w.ResponseWriter.(http.Hijacker)
	if !ok {
		return nil, nil, http.ErrNotSupported
	}
	return h.Hijack()
}

func (w *monitoringResponseWriter) Push(target string, opts *http.PushOptions) error {
	p, ok := w.ResponseWriter.(http.Pusher)
	if !ok {
		return http.ErrNotSupported
	}
	return p.Push(target, opts)
}

func (w *monitoringResponseWriter) ReadFrom(r io.Reader) (int64, error) {
	if rf, ok := w.ResponseWriter.(io.ReaderFrom); ok {
		return rf.ReadFrom(r)
	}
	return io.Copy(w.ResponseWriter, r)
}

func normalizeMetricPath(path string) string {
	path = strings.TrimSpace(path)
	if path == "" {
		return "/"
	}
	parts := strings.Split(path, "/")
	for i, part := range parts {
		if part == "" {
			continue
		}
		if isMetricDynamicSegment(part) {
			parts[i] = "{id}"
		}
	}
	return strings.Join(parts, "/")
}

func isMetricDynamicSegment(segment string) bool {
	if strings.HasPrefix(segment, "usr") || strings.HasPrefix(segment, "grp") || strings.HasPrefix(segment, "p2p") || strings.HasPrefix(segment, "chn") {
		return true
	}
	if _, err := strconv.ParseInt(segment, 10, 64); err == nil {
		return true
	}
	if uuidLikeMetricSegment.MatchString(segment) {
		return true
	}
	return false
}