From 2005f3ece1a479ebb9922a7195f836cc342ef3c7 Mon Sep 17 00:00:00 2001
From: Yohann D'ANELLO <yohann.danello@gmail.com>
Date: Tue, 29 Sep 2020 18:03:28 +0200
Subject: [PATCH] Store connected viewers in Prometheus and serve this amount

---
 go.mod                            |  2 ++
 go.sum                            |  1 +
 internal/monitoring/monitoring.go | 16 ++++++++++++++++
 stream/webrtc/webrtc.go           |  3 +++
 web/handler.go                    |  8 ++++++++
 web/web.go                        |  1 +
 6 files changed, 31 insertions(+)

diff --git a/go.mod b/go.mod
index 4fa3bbb..95daab6 100644
--- a/go.mod
+++ b/go.mod
@@ -5,10 +5,12 @@ go 1.13
 require (
 	github.com/go-ldap/ldap/v3 v3.2.3
 	github.com/haivision/srtgo v0.0.0-20200731151239-e00427ae473a
+	github.com/hashicorp/go.net v0.0.1
 	github.com/markbates/pkger v0.17.1
 	github.com/pion/rtp v1.6.0
 	github.com/pion/webrtc/v3 v3.0.0-beta.5
 	github.com/prometheus/client_golang v1.7.1
+	github.com/prometheus/client_model v0.2.0
 	github.com/spf13/viper v1.7.1
 	golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9
 )
diff --git a/go.sum b/go.sum
index a71c144..31cb6c3 100644
--- a/go.sum
+++ b/go.sum
@@ -148,6 +148,7 @@ github.com/hashicorp/go-sockaddr v1.0.0/go.mod h1:7Xibr9yA9JjQq1JpNB2Vw7kxv8xerX
 github.com/hashicorp/go-syslog v1.0.0/go.mod h1:qPfqrKkXGihmCqbJM2mZgkZGvKG1dFdvsLplgctolz4=
 github.com/hashicorp/go-uuid v1.0.0/go.mod h1:6SBZvOh/SIDV7/2o3Jml5SYk/TvGqwFJ/bN7x4byOro=
 github.com/hashicorp/go-uuid v1.0.1/go.mod h1:6SBZvOh/SIDV7/2o3Jml5SYk/TvGqwFJ/bN7x4byOro=
+github.com/hashicorp/go.net v0.0.1 h1:sNCoNyDEvN1xa+X0baata4RdcpKwcMS6DH+xwfqPgjw=
 github.com/hashicorp/go.net v0.0.1/go.mod h1:hjKkEWcCURg++eb33jQU7oqQcI9XDCnUzHA0oac0k90=
 github.com/hashicorp/golang-lru v0.5.0/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8=
 github.com/hashicorp/golang-lru v0.5.1/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8=
diff --git a/internal/monitoring/monitoring.go b/internal/monitoring/monitoring.go
index 2a10c50..8445f2d 100644
--- a/internal/monitoring/monitoring.go
+++ b/internal/monitoring/monitoring.go
@@ -7,6 +7,7 @@ import (
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
+	dto "github.com/prometheus/client_model/go"
 )
 
 // Options holds web package configuration
@@ -26,8 +27,23 @@ var (
 		Name: "ghostream_web_sessions_total",
 		Help: "The total amount of WebRTC sessions exchanged",
 	})
+
+	// WebRTCConnectedSessions is the total amount of WebRTC session exchange
+	WebRTCConnectedSessions = promauto.NewGauge(prometheus.GaugeOpts{
+		Name: "ghostream_webrtc_connected_sessions",
+		Help: "The current amount of opened WebRTC sessions",
+	})
 )
 
+func GetGaugeValue(metric prometheus.Gauge) float64 {
+	var m = &dto.Metric{}
+	if err := metric.Write(m); err != nil {
+		log.Fatal(err)
+		return 0
+	}
+	return m.Gauge.GetValue()
+}
+
 // Serve monitoring server that expose prometheus metrics
 func Serve(cfg *Options) {
 	mux := http.NewServeMux()
diff --git a/stream/webrtc/webrtc.go b/stream/webrtc/webrtc.go
index 5f7c346..ed23876 100644
--- a/stream/webrtc/webrtc.go
+++ b/stream/webrtc/webrtc.go
@@ -2,6 +2,7 @@ package webrtc
 
 import (
 	"fmt"
+	"gitlab.crans.org/nounous/ghostream/internal/monitoring"
 	"io"
 	"log"
 	"math/rand"
@@ -128,10 +129,12 @@ func newPeerHandler(remoteSdp webrtc.SessionDescription, cfg *Options) webrtc.Se
 			// Register tracks
 			videoTracks = append(videoTracks, videoTrack)
 			audioTracks = append(audioTracks, audioTrack)
+			monitoring.WebRTCConnectedSessions.Inc()
 		} else if connectionState == webrtc.ICEConnectionStateDisconnected {
 			// Unregister tracks
 			videoTracks = removeTrack(videoTracks, videoTrack)
 			audioTracks = removeTrack(audioTracks, audioTrack)
+			monitoring.WebRTCConnectedSessions.Dec()
 		}
 	})
 
diff --git a/web/handler.go b/web/handler.go
index f72ab8e..fef5ffc 100644
--- a/web/handler.go
+++ b/web/handler.go
@@ -84,3 +84,11 @@ func staticHandler() http.Handler {
 	staticFs := http.FileServer(pkger.Dir("/web/static"))
 	return http.StripPrefix("/static/", staticFs)
 }
+
+func statisticsHandler(w http.ResponseWriter, r *http.Request) {
+	// Display connected users stats
+	enc := json.NewEncoder(w)
+	enc.Encode(struct {
+		ConnectedViewers int
+	}{int(monitoring.GetGaugeValue(monitoring.WebRTCConnectedSessions))})
+}
diff --git a/web/web.go b/web/web.go
index cd44719..eca6e54 100644
--- a/web/web.go
+++ b/web/web.go
@@ -84,6 +84,7 @@ func Serve(rSdpChan chan webrtc.SessionDescription, lSdpChan chan webrtc.Session
 	mux := http.NewServeMux()
 	mux.HandleFunc("/", viewerHandler)
 	mux.Handle("/static/", staticHandler())
+	mux.HandleFunc("/_stats/", statisticsHandler)
 	log.Printf("HTTP server listening on %s", cfg.ListenAddress)
 	log.Fatal(http.ListenAndServe(cfg.ListenAddress, mux))
 }