Add prometheus endpoint

Author: tvrg

config/config.go

@@ -48,6 +48,7 @@ type Config struct {
 	AuthMode           string   `default:"turn" split_words:"true"`
 	CorsAllowedOrigins []string `split_words:"true"`
 	UsersFile          string   `split_words:"true"`
+	Prometheus         bool     `split_words:"true"`
 
 	CheckOrigin func(string) bool `ignored:"true" json:"-"`
 }

docs/config.md

@@ -69,4 +69,8 @@ SCREEGO_USERS_FILE=
 
 # The loglevel (one of: debug, info, warn, error)
 SCREEGO_LOG_LEVEL=info
-```
+
+# If screego should expose a prometheus endpoint at /metrics. The endpoint
+# requires basic authentication from a user in the users file.
+SCREEGO_PROMETHEUS=false
+```

go.mod

@@ -15,6 +15,7 @@ require (
 	github.com/phayes/freeport v0.0.0-20180830031419-95f893ade6f2
 	github.com/pion/transport v0.10.1 // indirect
 	github.com/pion/turn/v2 v2.0.4
+	github.com/prometheus/client_golang v1.7.1
 	github.com/rs/xid v1.2.1
 	github.com/rs/zerolog v1.19.0
 	github.com/stretchr/testify v1.6.1

go.sum

@@ -3,12 +3,13 @@ package router
 import (
 	"encoding/json"
 	"github.com/rs/zerolog/hlog"
-	"github.com/rs/zerolog/log"
 	"net/http"
 	"time"
 
 	"github.com/gorilla/handlers"
 	"github.com/gorilla/mux"
+	"github.com/prometheus/client_golang/prometheus/promhttp"
+	"github.com/rs/zerolog/log"
 	"github.com/screego/server/auth"
 	"github.com/screego/server/config"
 	"github.com/screego/server/ui"
@@ -42,6 +43,10 @@ func Router(conf config.Config, rooms *ws.Rooms, users *auth.Users, version stri
 			Version:  version,
 		})
 	})
+	if conf.Prometheus {
+		log.Info().Msg("Prometheus enabled")
+		router.Methods("GET").Path("/metrics").Handler(basicAuth(promhttp.Handler(), users))
+	}
 
 	ui.Register(router)
 
@@ -57,3 +62,19 @@ func accessLogger(r *http.Request, status, size int, dur time.Duration) {
 		Str("duration", dur.String()).
 		Msg("HTTP")
 }
+
+func basicAuth(handler http.Handler, users *auth.Users) http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+
+		user, pass, ok := r.BasicAuth()
+
+		if !ok || !users.Validate(user, pass) {
+			w.Header().Set("WWW-Authenticate", `Basic realm="screego"`)
+			w.WriteHeader(401)
+			_, _ = w.Write([]byte("Unauthorised.\n"))
+			return
+		}
+
+		handler.ServeHTTP(w, r)
+	}
+}

router/router.go

@@ -50,3 +50,7 @@ SCREEGO_USERS_FILE=
 
 # The loglevel (one of: debug, info, warn, error)
 SCREEGO_LOG_LEVEL=info
+
+# If screego should expose a prometheus endpoint at /metrics. The endpoint
+# requires basic authentication from a user in the users file.
+SCREEGO_PROMETHEUS=false

screego.config.example

@@ -0,0 +1,21 @@
+package ws
+
+import (
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promauto"
+)
+
+var (
+	roomGauge = promauto.NewGauge(prometheus.GaugeOpts{
+		Name: "screego_room_total",
+		Help: "The total number of rooms",
+	})
+	userGauge = promauto.NewGauge(prometheus.GaugeOpts{
+		Name: "screego_user_total",
+		Help: "The total number of users",
+	})
+	sessionGauge = promauto.NewGauge(prometheus.GaugeOpts{
+		Name: "screego_session_total",
+		Help: "The total number of sessions",
+	})
+)

ws/prometheus.go

@@ -65,10 +65,26 @@ func (r *Rooms) Upgrade(w http.ResponseWriter, req *http.Request) {
 }
 
 func (r *Rooms) Start() {
+	prometheusTicker := time.NewTicker(time.Second * 5)
+	if !r.config.Prometheus {
+		prometheusTicker.Stop()
+	}
 	for {
-		msg := <-r.Incoming
-		if err := msg.Incoming.Execute(r, msg.Info); err != nil {
-			msg.Info.Close <- err.Error()
+		select {
+		case msg := <-r.Incoming:
+			if err := msg.Incoming.Execute(r, msg.Info); err != nil {
+				msg.Info.Close <- err.Error()
+			}
+		case <-prometheusTicker.C:
+			roomGauge.Set(float64(len(r.Rooms)))
+			users := 0
+			sessions := 0
+			for _, room := range r.Rooms {
+				users += len(room.Users)
+				sessions += len(room.Sessions)
+			}
+			userGauge.Set(float64(users))
+			sessionGauge.Set(float64(sessions))
 		}
 	}
 }