diff --git a/cmd/agent.go b/cmd/agent.go index 9b38e1a..e33dc6a 100644 --- a/cmd/agent.go +++ b/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 { diff --git a/cmd/root.go b/cmd/root.go index b5b25d4..9669770 100644 --- a/cmd/root.go +++ b/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) + } + }() +} diff --git a/go.mod b/go.mod index 941f5fe..48ebbdd 100644 --- a/go.mod +++ b/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