feat: provide a TrafficCounter interface on server, remove prometheus client dependency from core

app/cmd/prom.go

@@ -0,0 +1,71 @@
+package main
+
+import (
+	"github.com/apernet/hysteria/core/cs"
+	"github.com/prometheus/client_golang/prometheus"
+)
+
+type prometheusTrafficCounter struct {
+	reg            *prometheus.Registry
+	upCounterVec   *prometheus.CounterVec
+	downCounterVec *prometheus.CounterVec
+	connGaugeVec   *prometheus.GaugeVec
+	counterMap     map[string]counters
+}
+
+type counters struct {
+	UpCounter   prometheus.Counter
+	DownCounter prometheus.Counter
+	ConnGauge   prometheus.Gauge
+}
+
+func NewPrometheusTrafficCounter(reg *prometheus.Registry) cs.TrafficCounter {
+	c := &prometheusTrafficCounter{
+		reg: reg,
+		upCounterVec: prometheus.NewCounterVec(prometheus.CounterOpts{
+			Name: "hysteria_traffic_uplink_bytes_total",
+		}, []string{"auth"}),
+		downCounterVec: prometheus.NewCounterVec(prometheus.CounterOpts{
+			Name: "hysteria_traffic_downlink_bytes_total",
+		}, []string{"auth"}),
+		connGaugeVec: prometheus.NewGaugeVec(prometheus.GaugeOpts{
+			Name: "hysteria_active_conn",
+		}, []string{"auth"}),
+		counterMap: make(map[string]counters),
+	}
+	reg.MustRegister(c.upCounterVec, c.downCounterVec, c.connGaugeVec)
+	return c
+}
+
+func (c *prometheusTrafficCounter) getCounters(auth string) counters {
+	cts, ok := c.counterMap[auth]
+	if !ok {
+		cts = counters{
+			UpCounter:   c.upCounterVec.WithLabelValues(auth),
+			DownCounter: c.downCounterVec.WithLabelValues(auth),
+			ConnGauge:   c.connGaugeVec.WithLabelValues(auth),
+		}
+		c.counterMap[auth] = cts
+	}
+	return cts
+}
+
+func (c *prometheusTrafficCounter) Rx(auth string, n int) {
+	cts := c.getCounters(auth)
+	cts.DownCounter.Add(float64(n))
+}
+
+func (c *prometheusTrafficCounter) Tx(auth string, n int) {
+	cts := c.getCounters(auth)
+	cts.UpCounter.Add(float64(n))
+}
+
+func (c *prometheusTrafficCounter) IncConn(auth string) {
+	cts := c.getCounters(auth)
+	cts.ConnGauge.Inc()
+}
+
+func (c *prometheusTrafficCounter) DecConn(auth string) {
+	cts := c.getCounters(auth)
+	cts.ConnGauge.Dec()
+}

app/cmd/server.go

@@ -192,9 +192,10 @@ func server(config *serverConfig) {
 		aclEngine.DefaultAction = acl.ActionDirect
 	}
 	// Prometheus
-	var promReg *prometheus.Registry
+	var trafficCounter cs.TrafficCounter
 	if len(config.PrometheusListen) > 0 {
-		promReg = prometheus.NewRegistry()
+		promReg := prometheus.NewRegistry()
+		trafficCounter = NewPrometheusTrafficCounter(promReg)
 		go func() {
 			http.Handle("/metrics", promhttp.HandlerFor(promReg, promhttp.HandlerOpts{}))
 			err := http.ListenAndServe(config.PrometheusListen, nil)
@@ -218,7 +219,7 @@ func server(config *serverConfig) {
 	up, down, _ := config.Speed()
 	server, err := cs.NewServer(tlsConfig, quicConfig, pktConn,
 		transport.DefaultServerTransport, up, down, config.DisableUDP, aclEngine,
-		connectFunc, disconnectFunc, tcpRequestFunc, tcpErrorFunc, udpRequestFunc, udpErrorFunc, promReg)
+		connectFunc, disconnectFunc, tcpRequestFunc, tcpErrorFunc, udpRequestFunc, udpErrorFunc, trafficCounter)
 	if err != nil {
 		logrus.WithField("error", err).Fatal("Failed to initialize server")
 	}