From 2477dc08ec02674ff4a1d6556e7c73014c85c917 Mon Sep 17 00:00:00 2001
From: zhangyifa <zhangyifa@lanjingren.com>
Date: Fri, 3 Jan 2020 15:52:41 +0800
Subject: [PATCH] implement API /metrics (#1357)

---
 cmd/frps/root.go               |  3 ++
 conf/frps_full.ini             |  3 ++
 go.mod                         |  1 +
 models/config/server_common.go |  7 +++++
 server/service.go              |  6 ++++
 server/stats/metrics.go        | 54 ++++++++++++++++++++++++++++++++++
 6 files changed, 74 insertions(+)
 create mode 100644 server/stats/metrics.go

diff --git a/cmd/frps/root.go b/cmd/frps/root.go
index ec175fe3..7b5ffa87 100644
--- a/cmd/frps/root.go
+++ b/cmd/frps/root.go
@@ -59,6 +59,7 @@ var (
 	allowPorts        string
 	maxPoolCount      int64
 	maxPortsPerClient int64
+	enableMetrics     bool
 )
 
 func init() {
@@ -86,6 +87,7 @@ func init() {
 	rootCmd.PersistentFlags().StringVarP(&subDomainHost, "subdomain_host", "", "", "subdomain host")
 	rootCmd.PersistentFlags().StringVarP(&allowPorts, "allow_ports", "", "", "allow ports")
 	rootCmd.PersistentFlags().Int64VarP(&maxPortsPerClient, "max_ports_per_client", "", 0, "max ports per client")
+	rootCmd.PersistentFlags().BoolVarP(&enableMetrics, "enable_metrics", "", false, "enable metrics api for prometheus")
 }
 
 var rootCmd = &cobra.Command{
@@ -173,6 +175,7 @@ func parseServerCommonCfgFromCmd() (cfg config.ServerCommonConf, err error) {
 	cfg.LogMaxDays = logMaxDays
 	cfg.Token = token
 	cfg.SubDomainHost = subDomainHost
+	cfg.EnableMetrics = enableMetrics
 	if len(allowPorts) > 0 {
 		// e.g. 1000-2000,2001,2002,3000-4000
 		ports, errRet := util.ParseRangeNumbers(allowPorts)
diff --git a/conf/frps_full.ini b/conf/frps_full.ini
index 030a3b3a..a7805a68 100644
--- a/conf/frps_full.ini
+++ b/conf/frps_full.ini
@@ -69,6 +69,9 @@ subdomain_host = frps.com
 # if tcp stream multiplexing is used, default is true
 tcp_mux = true
 
+# if true, a /metrics API will be exposed for Prometheus
+enable_metrics = false
+
 # custom 404 page for HTTP requests
 # custom_404_page = /path/to/404.html
 
diff --git a/go.mod b/go.mod
index a71a97c0..48fc293b 100644
--- a/go.mod
+++ b/go.mod
@@ -17,6 +17,7 @@ require (
 	github.com/mattn/go-runewidth v0.0.4 // indirect
 	github.com/pires/go-proxyproto v0.0.0-20190111085350-4d51b51e3bfc
 	github.com/pkg/errors v0.8.0 // indirect
+	github.com/prometheus/client_golang v1.3.0
 	github.com/rakyll/statik v0.1.1
 	github.com/rodaine/table v1.0.0
 	github.com/spf13/cobra v0.0.3
diff --git a/models/config/server_common.go b/models/config/server_common.go
index df6b7a10..c7b27e41 100644
--- a/models/config/server_common.go
+++ b/models/config/server_common.go
@@ -135,6 +135,8 @@ type ServerCommonConf struct {
 	// UserConnTimeout specifies the maximum time to wait for a work
 	// connection. By default, this value is 10.
 	UserConnTimeout int64 `json:"user_conn_timeout"`
+	// if true, a /metrics API will be exposed for Prometheus
+	EnableMetrics bool `json:"enable_metrics"`
 	// HTTPPlugins specify the server plugins support HTTP protocol.
 	HTTPPlugins map[string]plugin.HTTPPluginOptions `json:"http_plugins"`
 }
@@ -170,6 +172,7 @@ func GetDefaultServerConf() ServerCommonConf {
 		HeartBeatTimeout:  90,
 		UserConnTimeout:   10,
 		Custom404Page:     "",
+		EnableMetrics:     false,
 		HTTPPlugins:       make(map[string]plugin.HTTPPluginOptions),
 	}
 }
@@ -369,6 +372,10 @@ func UnmarshalServerConfFromIni(content string) (cfg ServerCommonConf, err error
 		cfg.Custom404Page = tmpStr
 	}
 
+	if tmpStr, ok = conf.Get("common", "enable_metrics"); ok && tmpStr == "true" {
+		cfg.EnableMetrics = true
+	}
+
 	if tmpStr, ok = conf.Get("common", "heartbeat_timeout"); ok {
 		v, errRet := strconv.ParseInt(tmpStr, 10, 64)
 		if errRet != nil {
diff --git a/server/service.go b/server/service.go
index 122555af..a025b3b3 100644
--- a/server/service.go
+++ b/server/service.go
@@ -253,6 +253,12 @@ func NewService(cfg config.ServerCommonConf) (svr *Service, err error) {
 	}
 
 	svr.statsCollector = stats.NewInternalCollector(statsEnable)
+
+	if cfg.EnableMetrics {
+		ms := stats.NewMetricsServer(svr.statsCollector)
+		go ms.Serve()
+		log.Info("Metrics server listen on %s", "8080")
+	}
 	return
 }
 
diff --git a/server/stats/metrics.go b/server/stats/metrics.go
new file mode 100644
index 00000000..41eae42d
--- /dev/null
+++ b/server/stats/metrics.go
@@ -0,0 +1,54 @@
+package stats
+
+import (
+	"net/http"
+
+	"time"
+
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promhttp"
+)
+
+// MetricsServer is a HTTP server wrapper
+type MetricsServer struct {
+	collector Collector
+}
+
+// NewMetricsServer creates a MetricsServer
+func NewMetricsServer(c Collector) *MetricsServer {
+	return &MetricsServer{
+		collector: c,
+	}
+}
+
+// Serve exposes Prometheus metrics data
+func (s *MetricsServer) Serve() {
+	http.Handle("/metrics", promhttp.Handler())
+
+	timestampCounter := prometheus.NewCounter(
+		prometheus.CounterOpts{
+			Namespace: "frps_info",
+			Name:      "timestamp",
+			Help:      "unix nanosec timestamp the data is collected",
+		})
+	clientCounts := prometheus.NewCounter(
+		prometheus.CounterOpts{
+			Namespace: "frps_info",
+			Name:      "client_counts",
+			Help:      "number of connected clients",
+		})
+	prometheus.MustRegister(timestampCounter)
+	prometheus.MustRegister(clientCounts)
+
+	go func() {
+		for {
+			stats := s.collector.GetServer()
+			timestampCounter.Add(float64(time.Now().UnixNano()))
+			clientCounts.Add(float64(stats.ClientCounts))
+			time.Sleep(time.Second)
+		}
+	}()
+
+	// FIXME load from conf
+	http.ListenAndServe(":8080", nil)
+}