im_server/server/monitoring_http.go

package main

import (
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"net/http"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"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,}$`)
	clusterQueueHighSince = make(map[string]time.Time)
	clusterQueueHealthMu  sync.Mutex
)

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"
	}
	if globals.cluster != nil {
		expected := len(globals.cluster.nodes)
		live := 0
		for _, node := range globals.cluster.nodes {
			node.lock.Lock()
			connected := node.connected
			node.lock.Unlock()
			if connected {
				live++
			}
		}
		checks["cluster"] = fmt.Sprintf("%d/%d peers connected", live, expected)
		if live != expected {
			status = http.StatusServiceUnavailable
		}
		if healthy, detail := clusterQueueHealthy(globals.cluster, time.Now()); !healthy {
			checks["cluster_queue"] = detail
			status = http.StatusServiceUnavailable
		} else {
			checks["cluster_queue"] = detail
		}
	}

	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 clusterQueueHealthy(cluster *Cluster, now time.Time) (bool, string) {
	clusterQueueHealthMu.Lock()
	defer clusterQueueHealthMu.Unlock()
	healthy := true
	detail := "ok"
	for name, node := range cluster.nodes {
		depth, capacity := 0, 0
		for _, sender := range node.p2mSenders {
			depth += len(sender)
			capacity += cap(sender)
		}
		if capacity == 0 || depth*100 < capacity*90 {
			delete(clusterQueueHighSince, name)
			continue
		}
		since, exists := clusterQueueHighSince[name]
		if !exists {
			clusterQueueHighSince[name] = now
			continue
		}
		if now.Sub(since) >= 30*time.Second {
			healthy = false
			detail = fmt.Sprintf("peer %s queue %d/%d above 90%% for %s", name, depth, capacity,
				now.Sub(since).Round(time.Second))
		}
	}
	return healthy, detail
}

func serveLivez(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.WriteHeader(http.StatusOK)
	if req.Method != http.MethodHead {
		_, _ = io.WriteString(wrt, "ok\n")
	}
}

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
}