Merge pull request #33 from benjamin99/feature/expose-prom-metrics

feat: add the support to export the metrics on the agent cmd

Author: rueian

cmd/agent.go

@@ -23,23 +23,32 @@ import (
 )
 
 var (
-	AgentListenAddr string
-	AgentRenice     int64
+	AgentListenAddr  string
+	AgentRenice      int64
+	ExportPrometheus bool
+	PrometheusAddr   string
 )
 
 func init() {
 	rootCmd.AddCommand(agent)
 	agent.Flags().StringVarP(&AgentListenAddr, "ListenAddr", "", ":10000", "the tcp address for agent server to listen")
 	agent.Flags().Int64VarP(&AgentRenice, "Renice", "", -10, "try renice the sink pg process")
+	agent.Flags().BoolVarP(&ExportPrometheus, "ExportPrometheus", "", false, "export the prometheus metrics or not")
+	agent.Flags().StringVarP(&PrometheusAddr, "PrometheusAddr", "", ":2112", "the tcp address for prometheus server to listen")
 }
 
 var agent = &cobra.Command{
 	Use:   "agent",
 	Short: "run as a agent accepting remote config",
 	RunE: func(cmd *cobra.Command, args []string) (err error) {
+		if ExportPrometheus {
+			startPrometheusServer(PrometheusAddr)
+		}
+
 		logrus.WithFields(logrus.Fields{
 			"AgentListenAddr": AgentListenAddr,
 		}).Info("starting agent")
+
 		agent := &Agent{}
 		return serveGRPC(&pb.Agent_ServiceDesc, AgentListenAddr, agent, func() {
 			if err := agent.cleanup(); err != nil {

cmd/root.go

@@ -11,6 +11,7 @@ import (
 	"strings"
 	"syscall"
 
+	"github.com/prometheus/client_golang/prometheus/promhttp"
 	"github.com/rueian/pgcapture/pkg/sink"
 	"github.com/rueian/pgcapture/pkg/source"
 	"github.com/sirupsen/logrus"
@@ -106,3 +107,24 @@ func trimSlot(topic string) string {
 	topic = strings.ReplaceAll(topic, "-", "_")
 	return topic
 }
+
+func startPrometheusServer(addr string) {
+	handler := promhttp.Handler()
+	server := &http.Server{
+		Addr: addr,
+		Handler: http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
+			if req.Method == http.MethodGet && req.URL.Path == "/metrics" {
+				handler.ServeHTTP(w, req)
+			} else {
+				http.NotFound(w, req)
+			}
+		}),
+	}
+
+	logrus.WithFields(logrus.Fields{"addr": addr}).Info("starting prometheus server")
+	go func() {
+		if err := server.ListenAndServe(); err != nil {
+			logrus.Error(err)
+		}
+	}()
+}

go.mod

@@ -14,6 +14,7 @@ require (
 	github.com/jackc/pgtype v1.7.0
 	github.com/jackc/pgx/v4 v4.10.1
 	github.com/pganalyze/pg_query_go/v2 v2.0.2
+	github.com/prometheus/client_golang v1.11.1
 	github.com/sirupsen/logrus v1.6.0
 	github.com/spf13/cobra v1.7.0
 	github.com/streamnative/pulsar-admin-go v0.1.0
@@ -48,7 +49,6 @@ require (
 	github.com/mtibben/percent v0.2.1 // indirect
 	github.com/pierrec/lz4 v2.0.5+incompatible // indirect
 	github.com/pkg/errors v0.9.1 // indirect
-	github.com/prometheus/client_golang v1.11.1 // indirect
 	github.com/prometheus/client_model v0.2.0 // indirect
 	github.com/prometheus/common v0.26.0 // indirect
 	github.com/prometheus/procfs v0.6.0 // indirect