Setup basic prom metrics endpoint (#624)

* Setup basic prom metrics endpoint
* Use default prom handler to expose go runtime metrics

pkg/api/api.go

@@ -16,6 +16,7 @@ import (
 	"github.com/sosedoff/pgweb/pkg/client"
 	"github.com/sosedoff/pgweb/pkg/command"
 	"github.com/sosedoff/pgweb/pkg/connection"
+	"github.com/sosedoff/pgweb/pkg/metrics"
 	"github.com/sosedoff/pgweb/pkg/shared"
 	"github.com/sosedoff/pgweb/static"
 )
@@ -278,6 +279,12 @@ func Disconnect(c *gin.Context) {
 		return
 	}
 
+	if command.Opts.Sessions {
+		result := DbSessions.Remove(getSessionId(c.Request))
+		successResponse(c, gin.H{"success": result})
+		return
+	}
+
 	conn := DB(c)
 	if conn == nil {
 		badRequest(c, errNotConnected)
@@ -492,6 +499,8 @@ func GetTableConstraints(c *gin.Context) {
 
 // HandleQuery runs the database query
 func HandleQuery(query string, c *gin.Context) {
+	metrics.IncrementQueriesCount()
+
 	rawQuery, err := base64.StdEncoding.DecodeString(desanitize64(query))
 	if err == nil {
 		query = string(rawQuery)

pkg/api/routes.go

@@ -4,6 +4,7 @@ import (
 	"github.com/gin-gonic/gin"
 
 	"github.com/sosedoff/pgweb/pkg/command"
+	"github.com/sosedoff/pgweb/pkg/metrics"
 )
 
 func SetupMiddlewares(group *gin.RouterGroup) {
@@ -53,3 +54,9 @@ func SetupRoutes(router *gin.Engine) {
 	api.GET("/bookmarks", GetBookmarks)
 	api.GET("/export", DataExport)
 }
+
+func SetupMetrics(engine *gin.Engine) {
+	if command.Opts.MetricsEnabled && command.Opts.MetricsAddr == "" {
+		engine.GET("/metrics", gin.WrapH(metrics.Handler()))
+	}
+}

pkg/api/session_manager.go

@@ -7,6 +7,7 @@ import (
 	"github.com/sirupsen/logrus"
 
 	"github.com/sosedoff/pgweb/pkg/client"
+	"github.com/sosedoff/pgweb/pkg/metrics"
 )
 
 type SessionManager struct {
@@ -58,8 +59,10 @@ func (m *SessionManager) Get(id string) *client.Client {
 
 func (m *SessionManager) Add(id string, conn *client.Client) {
 	m.mu.Lock()
+	defer m.mu.Unlock()
+
 	m.sessions[id] = conn
-	m.mu.Unlock()
+	metrics.SetSessionsCount(len(m.sessions))
 }
 
 func (m *SessionManager) Remove(id string) bool {
@@ -72,6 +75,7 @@ func (m *SessionManager) Remove(id string) bool {
 		delete(m.sessions, id)
 	}
 
+	metrics.SetSessionsCount(len(m.sessions))
 	return ok
 }
 

pkg/cli/cli.go

@@ -19,6 +19,7 @@ import (
 	"github.com/sosedoff/pgweb/pkg/client"
 	"github.com/sosedoff/pgweb/pkg/command"
 	"github.com/sosedoff/pgweb/pkg/connection"
+	"github.com/sosedoff/pgweb/pkg/metrics"
 	"github.com/sosedoff/pgweb/pkg/util"
 )
 
@@ -199,9 +200,12 @@ func startServer() {
 
 	api.SetLogger(logger)
 	api.SetupRoutes(router)
+	api.SetupMetrics(router)
 
 	fmt.Println("Starting server...")
 	go func() {
+		metrics.SetHealty(true)
+
 		err := router.Run(fmt.Sprintf("%v:%v", options.HTTPHost, options.HTTPPort))
 		if err != nil {
 			fmt.Println("Cant start server:", err)
@@ -213,6 +217,18 @@ func startServer() {
 	}()
 }
 
+func startMetricsServer() {
+	serverAddr := fmt.Sprintf("%v:%v", command.Opts.HTTPHost, command.Opts.HTTPPort)
+	if options.MetricsAddr == serverAddr {
+		return
+	}
+
+	err := metrics.StartServer(logger, options.MetricsPath, options.MetricsAddr)
+	if err != nil {
+		logger.WithError(err).Fatal("unable to start prometheus metrics server")
+	}
+}
+
 func handleSignals() {
 	c := make(chan os.Signal, 1)
 	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
@@ -265,6 +281,10 @@ func Run() {
 		}
 	}
 
+	if options.MetricsEnabled && options.MetricsAddr != "" {
+		go startMetricsServer()
+	}
+
 	startServer()
 	openPage()
 	handleSignals()

pkg/command/options.go

@@ -58,6 +58,9 @@ type Options struct {
 	Cors                         bool   `long:"cors" description:"Enable Cross-Origin Resource Sharing (CORS)"`
 	CorsOrigin                   string `long:"cors-origin" description:"Allowed CORS origins" default:"*"`
 	BinaryCodec                  string `long:"binary-codec" description:"Codec for binary data serialization, one of 'none', 'hex', 'base58', 'base64'" default:"none"`
+	MetricsEnabled               bool   `long:"metrics" description:"Enable Prometheus metrics endpoint"`
+	MetricsPath                  string `long:"metrics-path" description:"Path prefix for Prometheus metrics endpoint" default:"/metrics"`
+	MetricsAddr                  string `long:"metrics-addr" description:"Listen host and port for Prometheus metrics server"`
 }
 
 var Opts Options

pkg/metrics/metrics.go

@@ -0,0 +1,39 @@
+package metrics
+
+import (
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promauto"
+)
+
+var (
+	sessionsGauge = promauto.NewGauge(prometheus.GaugeOpts{
+		Name: "pgweb_sessions_count",
+		Help: "Total number of database sessions",
+	})
+
+	queriesCounter = promauto.NewCounter(prometheus.CounterOpts{
+		Name: "pgweb_queries_count",
+		Help: "Total number of custom queries executed",
+	})
+
+	healtyGauge = promauto.NewGauge(prometheus.GaugeOpts{
+		Name: "pgweb_healty",
+		Help: "Server health status",
+	})
+)
+
+func IncrementQueriesCount() {
+	queriesCounter.Inc()
+}
+
+func SetSessionsCount(val int) {
+	sessionsGauge.Set(float64(val))
+}
+
+func SetHealty(val bool) {
+	healthy := 0.0
+	if val {
+		healthy = 1.0
+	}
+	healtyGauge.Set(float64(healthy))
+}

pkg/metrics/server.go

@@ -0,0 +1,19 @@
+package metrics
+
+import (
+	"net/http"
+
+	"github.com/prometheus/client_golang/prometheus/promhttp"
+	"github.com/sirupsen/logrus"
+)
+
+func Handler() http.Handler {
+	return promhttp.Handler()
+}
+
+func StartServer(logger *logrus.Logger, path string, addr string) error {
+	logger.WithField("addr", addr).WithField("path", path).Info("starting prometheus metrics server")
+
+	http.Handle(path, Handler())
+	return http.ListenAndServe(addr, nil)
+}