initial prometheus metrics support, simplfied some code

internal/api/v1/stats.go

@@ -10,7 +10,6 @@ import (
 	U "github.com/yusing/go-proxy/internal/api/v1/utils"
 	"github.com/yusing/go-proxy/internal/common"
 	"github.com/yusing/go-proxy/internal/config"
-	"github.com/yusing/go-proxy/internal/server"
 	"github.com/yusing/go-proxy/internal/utils/strutils"
 )
 
@@ -33,7 +32,6 @@ func StatsWS(w http.ResponseWriter, r *http.Request) {
 		}
 		originPats = append(originPats, localAddresses...)
 	}
-	U.LogInfo(r).Msgf("websocket API request from origins: %s", originPats)
 	if common.IsDebug {
 		originPats = []string{"*"}
 	}
@@ -62,9 +60,11 @@ func StatsWS(w http.ResponseWriter, r *http.Request) {
 	}
 }
 
+var startTime = time.Now()
+
 func getStats() map[string]any {
 	return map[string]any{
 		"proxies": config.Statistics(),
-		"uptime":  strutils.FormatDuration(server.GetProxyServer().Uptime()),
+		"uptime":  strutils.FormatDuration(time.Since(startTime)),
 	}
 }

internal/common/env.go

@@ -19,6 +19,7 @@ var (
 	IsDebug            = GetEnvBool("DEBUG", IsTest)
 	IsDebugSkipAuth    = GetEnvBool("DEBUG_SKIP_AUTH", false)
 	IsTrace            = GetEnvBool("TRACE", false) && IsDebug
+	IsProduction       = !IsTest && !IsDebug
 
 	ProxyHTTPAddr,
 	ProxyHTTPHost,
@@ -35,6 +36,12 @@ var (
 	APIHTTPPort,
 	APIHTTPURL = GetAddrEnv("API_ADDR", "127.0.0.1:8888", "http")
 
+	MetricsHTTPAddr,
+	MetricsHTTPHost,
+	MetricsHTTPPort,
+	MetricsHTTPURL = GetAddrEnv("PROMETHEUS_ADDR", "", "http")
+	PrometheusEnabled = MetricsHTTPURL != ""
+
 	APIJWTSecret    = decodeJWTKey(GetEnvString("API_JWT_SECRET", ""))
 	APIJWTTokenTTL  = GetDurationEnv("API_JWT_TOKEN_TTL", time.Hour)
 	APIUser         = GetEnvString("API_USER", "admin")
@@ -79,6 +86,9 @@ func GetEnvBool(key string, defaultValue bool) bool {
 
 func GetAddrEnv(key, defaultValue, scheme string) (addr, host, port, fullURL string) {
 	addr = GetEnvString(key, defaultValue)
+	if addr == "" {
+		return
+	}
 	host, port, err := net.SplitHostPort(addr)
 	if err != nil {
 		log.Fatal().Msgf("env %s: invalid address: %s", key, addr)

internal/config/query.go

@@ -9,7 +9,6 @@ import (
 	"github.com/yusing/go-proxy/internal/proxy/entry"
 	"github.com/yusing/go-proxy/internal/route"
 	proxy "github.com/yusing/go-proxy/internal/route/provider"
-	F "github.com/yusing/go-proxy/internal/utils/functional"
 	"github.com/yusing/go-proxy/internal/utils/strutils"
 )
 
@@ -139,13 +138,12 @@ func Statistics() map[string]any {
 	providerStats := make(map[string]proxy.ProviderStats)
 
 	instance.providers.RangeAll(func(name string, p *proxy.Provider) {
-		providerStats[name] = p.Statistics()
-	})
+		stats := p.Statistics()
+		providerStats[name] = stats
 
-	for _, stats := range providerStats {
 		nTotalRPs += stats.NumRPs
 		nTotalStreams += stats.NumStreams
-	}
+	})
 
 	return map[string]any{
 		"num_total_streams":         nTotalStreams,
@@ -153,14 +151,3 @@ func Statistics() map[string]any {
 		"providers":                 providerStats,
 	}
 }
-
-func FindRoute(alias string) *route.Route {
-	return F.MapFind(instance.providers,
-		func(p *proxy.Provider) (*route.Route, bool) {
-			if route, ok := p.GetRoute(alias); ok {
-				return route, true
-			}
-			return nil, false
-		},
-	)
-}

internal/docker/idlewatcher/types/waker.go

@@ -11,4 +11,5 @@ type Waker interface {
 	health.HealthMonitor
 	http.Handler
 	net.Stream
+	Wake() error
 }

internal/docker/idlewatcher/watcher.go

@@ -98,6 +98,10 @@ func registerWatcher(providerSubtask task.Task, entry entry.Entry, waker *waker)
 	return w, nil
 }
 
+func (w *Watcher) Wake() error {
+	return w.wakeIfStopped()
+}
+
 // WakeDebug logs a debug message related to waking the container.
 func (w *Watcher) WakeDebug() *zerolog.Event {
 	return w.Debug().Str("action", "wake")

internal/metrics/http_handler.go

@@ -0,0 +1,13 @@
+package metrics
+
+import (
+	"net/http"
+
+	"github.com/prometheus/client_golang/prometheus/promhttp"
+)
+
+func NewHandler() http.Handler {
+	mux := http.NewServeMux()
+	mux.Handle("/metrics", promhttp.Handler())
+	return mux
+}

internal/metrics/http_metrics.go

@@ -0,0 +1,82 @@
+package metrics
+
+import (
+	"strings"
+
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/yusing/go-proxy/internal/common"
+)
+
+type (
+	RouteMetrics struct {
+		HTTPReqTotal,
+		HTTP2xx3xx,
+		HTTP4xx,
+		HTTP5xx *Counter
+		HTTPReqElapsed *Gauge
+	}
+	HTTPRouteMetricLabels struct {
+		Service, Method, Host, Visitor, Path string
+	}
+)
+
+var rm RouteMetrics
+
+const (
+	routerNamespace     = "router"
+	routerHTTPSubsystem = "http"
+)
+
+func GetRouteMetrics() *RouteMetrics {
+	return &rm
+}
+
+func (lbl HTTPRouteMetricLabels) toPromLabels() prometheus.Labels {
+	return prometheus.Labels{
+		"service": lbl.Service,
+		"method":  lbl.Method,
+		"host":    lbl.Host,
+		"visitor": lbl.Visitor,
+		"path":    lbl.Path,
+	}
+}
+
+func init() {
+	if !common.PrometheusEnabled {
+		return
+	}
+	lbls := []string{"service", "method", "host", "visitor", "path"}
+	partitionsHelp := ", partitioned by " + strings.Join(lbls, ", ")
+	rm = RouteMetrics{
+		HTTPReqTotal: NewCounter(prometheus.CounterOpts{
+			Namespace: routerNamespace,
+			Subsystem: routerHTTPSubsystem,
+			Name:      "req_total",
+			Help:      "How many requests processed" + partitionsHelp,
+		}),
+		HTTP2xx3xx: NewCounter(prometheus.CounterOpts{
+			Namespace: routerNamespace,
+			Subsystem: routerHTTPSubsystem,
+			Name:      "req_ok_count",
+			Help:      "How many 2xx-3xx requests processed" + partitionsHelp,
+		}, lbls...),
+		HTTP4xx: NewCounter(prometheus.CounterOpts{
+			Namespace: routerNamespace,
+			Subsystem: routerHTTPSubsystem,
+			Name:      "req_4xx_count",
+			Help:      "How many 4xx requests processed" + partitionsHelp,
+		}, lbls...),
+		HTTP5xx: NewCounter(prometheus.CounterOpts{
+			Namespace: routerNamespace,
+			Subsystem: routerHTTPSubsystem,
+			Name:      "req_5xx_count",
+			Help:      "How many 5xx requests processed" + partitionsHelp,
+		}, lbls...),
+		HTTPReqElapsed: NewGauge(prometheus.GaugeOpts{
+			Namespace: routerNamespace,
+			Subsystem: routerHTTPSubsystem,
+			Name:      "req_elapsed_ms",
+			Help:      "How long it took to process the request" + partitionsHelp,
+		}, lbls...),
+	}
+}

internal/metrics/metric.go

@@ -0,0 +1,73 @@
+package metrics
+
+import "github.com/prometheus/client_golang/prometheus"
+
+type (
+	Counter struct {
+		collector prometheus.Counter
+		mv        *prometheus.CounterVec
+	}
+	Gauge struct {
+		collector prometheus.Gauge
+		mv        *prometheus.GaugeVec
+	}
+	Labels interface {
+		toPromLabels() prometheus.Labels
+	}
+)
+
+func NewCounter(opts prometheus.CounterOpts, labels ...string) *Counter {
+	m := &Counter{
+		mv: prometheus.NewCounterVec(opts, labels),
+	}
+	if len(labels) == 0 {
+		m.collector = m.mv.WithLabelValues()
+		m.collector.Add(0)
+	}
+	prometheus.MustRegister(m)
+	return m
+}
+
+func NewGauge(opts prometheus.GaugeOpts, labels ...string) *Gauge {
+	m := &Gauge{
+		mv: prometheus.NewGaugeVec(opts, labels),
+	}
+	if len(labels) == 0 {
+		m.collector = m.mv.WithLabelValues()
+		m.collector.Set(0)
+	}
+	prometheus.MustRegister(m)
+	return m
+}
+
+func (c *Counter) Collect(ch chan<- prometheus.Metric) {
+	c.mv.Collect(ch)
+}
+
+func (c *Counter) Describe(ch chan<- *prometheus.Desc) {
+	c.mv.Describe(ch)
+}
+
+func (c *Counter) Inc() {
+	c.collector.Inc()
+}
+
+func (c *Counter) With(l Labels) prometheus.Counter {
+	return c.mv.With(l.toPromLabels())
+}
+
+func (g *Gauge) Collect(ch chan<- prometheus.Metric) {
+	g.mv.Collect(ch)
+}
+
+func (g *Gauge) Describe(ch chan<- *prometheus.Desc) {
+	g.mv.Describe(ch)
+}
+
+func (g *Gauge) Set(v float64) {
+	g.collector.Set(v)
+}
+
+func (g *Gauge) With(l Labels) prometheus.Gauge {
+	return g.mv.With(l.toPromLabels())
+}

internal/metrics/router_metrics.go

@@ -0,0 +1,20 @@
+package metrics
+
+import "github.com/prometheus/client_golang/prometheus"
+
+func InitRouterMetrics(getRPsCount func() int, getStreamsCount func() int) {
+	prometheus.MustRegister(prometheus.NewGaugeFunc(prometheus.GaugeOpts{
+		Namespace: "entrypoint",
+		Name:      "num_reverse_proxies",
+		Help:      "The number of reverse proxies",
+	}, func() float64 {
+		return float64(getRPsCount())
+	}))
+	prometheus.MustRegister(prometheus.NewGaugeFunc(prometheus.GaugeOpts{
+		Namespace: "entrypoint",
+		Name:      "num_streams",
+		Help:      "The number of streams",
+	}, func() float64 {
+		return float64(getStreamsCount())
+	}))
+}

internal/net/http/dummy_response_writer.go

@@ -1,15 +0,0 @@
-package http
-
-import "net/http"
-
-type DummyResponseWriter struct{}
-
-func (w DummyResponseWriter) Header() http.Header {
-	return make(http.Header)
-}
-
-func (w DummyResponseWriter) Write([]byte) (_ int, _ error) {
-	return
-}
-
-func (w DummyResponseWriter) WriteHeader(int) {}

internal/net/http/loadbalancer/loadbalancer.go

@@ -1,7 +1,6 @@
 package loadbalancer
 
 import (
-	"context"
 	"net/http"
 	"sync"
 	"time"
@@ -10,7 +9,6 @@ import (
 	"github.com/yusing/go-proxy/internal/common"
 	idlewatcher "github.com/yusing/go-proxy/internal/docker/idlewatcher/types"
 	E "github.com/yusing/go-proxy/internal/error"
-	gphttp "github.com/yusing/go-proxy/internal/net/http"
 	"github.com/yusing/go-proxy/internal/net/http/middleware"
 	"github.com/yusing/go-proxy/internal/task"
 	"github.com/yusing/go-proxy/internal/watcher/health"
@@ -225,18 +223,15 @@ func (lb *LoadBalancer) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
 		return
 	}
 	if r.Header.Get(common.HeaderCheckRedirect) != "" {
-		ctx, cancel := context.WithTimeout(r.Context(), 1*time.Second)
-		defer cancel()
-		// send dummy request to wake all servers
-		var dummyRW gphttp.DummyResponseWriter
+		// wake all servers
 		for _, srv := range srvs {
 			// wake only if server implements Waker
-			_, ok := srv.handler.(idlewatcher.Waker)
-			if !ok {
-				continue
+			waker, ok := srv.handler.(idlewatcher.Waker)
+			if ok {
+				if err := waker.Wake(); err != nil {
+					lb.Err(err).Msgf("failed to wake server %s", srv.Name)
+				}
 			}
-			wakeReq := r.Clone(ctx)
-			srv.ServeHTTP(dummyRW, wakeReq)
 		}
 	}
 	lb.impl.ServeHTTP(srvs, rw, r)

internal/net/http/middleware/middleware.go

@@ -157,8 +157,8 @@ func patchReverseProxy(rpName string, rp *ReverseProxy, middlewares []*Middlewar
 	mid := BuildMiddlewareFromChain(rpName, middlewares)
 
 	if mid.before != nil {
-		ori := rp.ServeHTTP
-		rp.ServeHTTP = func(w http.ResponseWriter, r *http.Request) {
+		ori := rp.HandlerFunc
+		rp.HandlerFunc = func(w http.ResponseWriter, r *http.Request) {
 			mid.before(ori, w, r)
 		}
 	}

internal/net/http/reverse_proxy_mod.go

@@ -10,6 +10,7 @@ package http
 // Copyright (c) 2024 yusing
 
 import (
+	"bufio"
 	"bytes"
 	"context"
 	"errors"
@@ -22,8 +23,11 @@ import (
 	"net/url"
 	"strings"
 	"sync"
+	"time"
 
 	"github.com/rs/zerolog"
+	"github.com/yusing/go-proxy/internal/common"
+	"github.com/yusing/go-proxy/internal/metrics"
 	"github.com/yusing/go-proxy/internal/net/types"
 	U "github.com/yusing/go-proxy/internal/utils"
 	"golang.org/x/net/http/httpguts"
@@ -86,12 +90,52 @@ type ReverseProxy struct {
 	// implementation is used.
 	ModifyResponse func(*http.Response) error
 
-	ServeHTTP http.HandlerFunc
+	HandlerFunc http.HandlerFunc
 
 	TargetName string
 	TargetURL  types.URL
 }
 
+type httpMetricLogger struct {
+	http.ResponseWriter
+	labels metrics.HTTPRouteMetricLabels
+}
+
+// WriteHeader implements http.ResponseWriter.
+func (l *httpMetricLogger) WriteHeader(status int) {
+	l.ResponseWriter.WriteHeader(status)
+	go func() {
+		m := metrics.GetRouteMetrics()
+		m.HTTPReqTotal.Inc()
+
+		// ignore 1xx
+		switch {
+		case status >= 500:
+			m.HTTP5xx.With(l.labels).Inc()
+		case status >= 400:
+			m.HTTP4xx.With(l.labels).Inc()
+		case status >= 200:
+			m.HTTP2xx3xx.With(l.labels).Inc()
+		}
+	}()
+}
+
+// Hijack hijacks the connection.
+func (l *httpMetricLogger) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	if h, ok := l.ResponseWriter.(http.Hijacker); ok {
+		return h.Hijack()
+	}
+
+	return nil, nil, fmt.Errorf("not a hijacker: %T", l.ResponseWriter)
+}
+
+// Flush sends any buffered data to the client.
+func (l *httpMetricLogger) Flush() {
+	if flusher, ok := l.ResponseWriter.(http.Flusher); ok {
+		flusher.Flush()
+	}
+}
+
 func singleJoiningSlash(a, b string) string {
 	aslash := strings.HasSuffix(a, "/")
 	bslash := strings.HasPrefix(b, "/")
@@ -157,7 +201,7 @@ func NewReverseProxy(name string, target types.URL, transport http.RoundTripper)
 		TargetName: name,
 		TargetURL:  target,
 	}
-	rp.ServeHTTP = rp.serveHTTP
+	rp.HandlerFunc = rp.handler
 	return rp
 }
 
@@ -225,9 +269,32 @@ func (p *ReverseProxy) modifyResponse(rw http.ResponseWriter, res *http.Response
 	return true
 }
 
-func (p *ReverseProxy) serveHTTP(rw http.ResponseWriter, req *http.Request) {
-	if _, ok := rw.(DummyResponseWriter); ok {
-		return
+func (p *ReverseProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
+	p.HandlerFunc(rw, req)
+}
+
+func (p *ReverseProxy) handler(rw http.ResponseWriter, req *http.Request) {
+	if common.PrometheusEnabled {
+		t := time.Now()
+		visitor, _, err := net.SplitHostPort(req.RemoteAddr)
+		if err != nil {
+			visitor = req.RemoteAddr
+		}
+		lbls := metrics.HTTPRouteMetricLabels{
+			Service: p.TargetName,
+			Method:  req.Method,
+			Host:    req.Host,
+			Visitor: visitor,
+			Path:    req.URL.Path,
+		}
+		rw = &httpMetricLogger{
+			ResponseWriter: rw,
+			labels:         lbls,
+		}
+		defer func() {
+			duration := time.Since(t)
+			metrics.GetRouteMetrics().HTTPReqElapsed.With(lbls).Set(float64(duration.Milliseconds()))
+		}()
 	}
 
 	transport := p.Transport

internal/route/http.go

@@ -38,10 +38,6 @@ type (
 	}
 
 	SubdomainKey = PT.Alias
-
-	ReverseProxyHandler struct {
-		*gphttp.ReverseProxy
-	}
 )
 
 var (
@@ -52,10 +48,6 @@ var (
 	// globalMux    = http.NewServeMux() // TODO: support regex subdomain matching.
 )
 
-func (rp ReverseProxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
-	rp.ReverseProxy.ServeHTTP(w, r)
-}
-
 func GetReverseProxies() F.Map[string, *HTTPRoute] {
 	return httpRoutes
 }
@@ -77,10 +69,11 @@ func NewHTTPRoute(entry *entry.ReverseProxyEntry) (impl, E.Error) {
 		trans = gphttp.DefaultTransport.Clone()
 	}
 
-	rp := gphttp.NewReverseProxy(string(entry.Alias), entry.URL, trans)
+	service := string(entry.Alias)
+	rp := gphttp.NewReverseProxy(service, entry.URL, trans)
 
 	if len(entry.Middlewares) > 0 {
-		err := middleware.PatchReverseProxy(string(entry.Alias), rp, entry.Middlewares)
+		err := middleware.PatchReverseProxy(service, rp, entry.Middlewares)
 		if err != nil {
 			return nil, err
 		}
@@ -136,11 +129,11 @@ func (r *HTTPRoute) Start(providerSubtask task.Task) E.Error {
 	if r.handler == nil {
 		switch {
 		case len(r.PathPatterns) == 1 && r.PathPatterns[0] == "/":
-			r.handler = ReverseProxyHandler{r.rp}
+			r.handler = r.rp
 		default:
 			mux := http.NewServeMux()
 			for _, p := range r.PathPatterns {
-				mux.HandleFunc(string(p), r.rp.ServeHTTP)
+				mux.HandleFunc(string(p), r.rp.HandlerFunc)
 			}
 			r.handler = mux
 		}

internal/route/metrics.go

@@ -0,0 +1,7 @@
+package route
+
+import "github.com/yusing/go-proxy/internal/metrics"
+
+func init() {
+	metrics.InitRouterMetrics(httpRoutes.Size, streamRoutes.Size)
+}

internal/server/instance.go

@@ -1,25 +0,0 @@
-package server
-
-var proxyServer, apiServer *Server
-
-func InitProxyServer(opt Options) *Server {
-	if proxyServer == nil {
-		proxyServer = NewServer(opt)
-	}
-	return proxyServer
-}
-
-func InitAPIServer(opt Options) *Server {
-	if apiServer == nil {
-		apiServer = NewServer(opt)
-	}
-	return apiServer
-}
-
-func GetProxyServer() *Server {
-	return proxyServer
-}
-
-func GetAPIServer() *Server {
-	return apiServer
-}

internal/server/server.go

@@ -38,6 +38,12 @@ type Options struct {
 	Handler         http.Handler
 }
 
+func StartServer(opt Options) (s *Server) {
+	s = NewServer(opt)
+	s.Start()
+	return s
+}
+
 func NewServer(opt Options) (s *Server) {
 	var httpSer, httpsSer *http.Server
 	var httpHandler http.Handler

internal/watcher/events/event_queue.go

@@ -50,8 +50,25 @@ func NewEventQueue(parent task.Task, flushInterval time.Duration, onFlush OnFlus
 }
 
 func (e *EventQueue) Start(eventCh <-chan Event, errCh <-chan E.Error) {
+	if common.IsProduction {
+		origOnFlush := e.onFlush
+		// recover panic in onFlush when in production mode
+		e.onFlush = func(flushTask task.Task, events []Event) {
+			defer func() {
+				if err := recover(); err != nil {
+					e.onError(E.New("recovered panic in onFlush").
+						Withf("%v", err).
+						Subject(e.task.Parent().String()))
+				}
+			}()
+			origOnFlush(flushTask, events)
+		}
+	}
+
 	go func() {
 		defer e.ticker.Stop()
+		defer e.task.Finish(nil)
+
 		for {
 			select {
 			case <-e.task.Context().Done():
@@ -61,18 +78,7 @@ func (e *EventQueue) Start(eventCh <-chan Event, errCh <-chan E.Error) {
 					flushTask := e.task.Subtask("flush events")
 					queue := e.queue
 					e.queue = make([]Event, 0, eventQueueCapacity)
-					if !common.IsDebug {
-						go func() {
-							defer func() {
-								if err := recover(); err != nil {
-									e.onError(E.Errorf("recovered panic in onFlush: %v", err).Subject(e.task.Parent().String()))
-								}
-							}()
-							e.onFlush(flushTask, queue)
-						}()
-					} else {
-						go e.onFlush(flushTask, queue)
-					}
+					go e.onFlush(flushTask, queue)
 					flushTask.Wait()
 				}
 				e.ticker.Reset(e.flushInterval)