From 3dae9b19991c6240fce93fe3d88dabd6f2ef5f89 Mon Sep 17 00:00:00 2001
From: Tulir Asokan <tulir@maunium.net>
Date: Sun, 27 Sep 2020 22:30:08 +0300
Subject: [PATCH] Add connection state to prometheus metrics

---
 commands.go     |  2 ++
 metrics.go      | 39 +++++++++++++++++++++++++++++++++++++++
 provisioning.go |  2 ++
 user.go         |  5 +++++
 4 files changed, 48 insertions(+)

diff --git a/commands.go b/commands.go
index a2ed2c8..8b2628e 100644
--- a/commands.go
+++ b/commands.go
@@ -551,6 +551,7 @@ func (handler *CommandHandler) CommandDeleteConnection(ce *CommandEvent) {
 	}
 	ce.User.Conn.RemoveHandlers()
 	ce.User.Conn = nil
+	ce.User.bridge.Metrics.TrackConnectionState(ce.User.JID, false)
 	ce.Reply("Successfully disconnected. Use the `reconnect` command to reconnect.")
 }
 
@@ -572,6 +573,7 @@ func (handler *CommandHandler) CommandDisconnect(ce *CommandEvent) {
 	} else if len(sess.Wid) > 0 {
 		ce.User.SetSession(&sess)
 	}
+	ce.User.bridge.Metrics.TrackConnectionState(ce.User.JID, false)
 	ce.Reply("Successfully disconnected. Use the `reconnect` command to reconnect.")
 }
 
diff --git a/metrics.go b/metrics.go
index 385de7a..edd27a7 100644
--- a/metrics.go
+++ b/metrics.go
@@ -30,6 +30,7 @@ import (
 	"maunium.net/go/mautrix/id"
 
 	"maunium.net/go/mautrix-whatsapp/database"
+	"maunium.net/go/mautrix-whatsapp/types"
 )
 
 type MetricsHandler struct {
@@ -52,6 +53,9 @@ type MetricsHandler struct {
 	encryptedPrivateCount   prometheus.Gauge
 	unencryptedGroupCount   prometheus.Gauge
 	unencryptedPrivateCount prometheus.Gauge
+
+	connected *prometheus.GaugeVec
+	loggedIn  *prometheus.GaugeVec
 }
 
 func NewMetricsHandler(address string, log log.Logger, db *database.Database) *MetricsHandler {
@@ -94,6 +98,15 @@ func NewMetricsHandler(address string, log log.Logger, db *database.Database) *M
 		encryptedPrivateCount:   portalCount.With(prometheus.Labels{"type": "private", "encrypted": "true"}),
 		unencryptedGroupCount:   portalCount.With(prometheus.Labels{"type": "group", "encrypted": "false"}),
 		unencryptedPrivateCount: portalCount.With(prometheus.Labels{"type": "private", "encrypted": "false"}),
+
+		loggedIn: promauto.NewGaugeVec(prometheus.GaugeOpts{
+			Name: "bridge_logged_in",
+			Help: "Users logged into the bridge",
+		}, []string{"jid", "bridge_logged_in"}),
+		connected: promauto.NewGaugeVec(prometheus.GaugeOpts{
+			Name: "bridge_connected",
+			Help: "Users connected to WhatsApp",
+		}, []string{"jid", "bridge_connected"}),
 	}
 }
 
@@ -119,6 +132,32 @@ func (mh *MetricsHandler) TrackDisconnection(userID id.UserID) {
 	mh.disconnections.With(prometheus.Labels{"user_id": string(userID)}).Inc()
 }
 
+func (mh *MetricsHandler) TrackLoginState(jid types.WhatsAppID, loggedIn bool) {
+	if !mh.running {
+		return
+	}
+	var loggedInVal float64 = 0
+	if loggedIn {
+		loggedInVal = 1
+	}
+	metric := mh.loggedIn.MustCurryWith(prometheus.Labels{"jid": jid})
+	metric.With(prometheus.Labels{"bridge_logged_in": "true"}).Set(loggedInVal)
+	metric.With(prometheus.Labels{"bridge_logged_in": "false"}).Set(1 - loggedInVal)
+}
+
+func (mh *MetricsHandler) TrackConnectionState(jid types.WhatsAppID, connected bool) {
+	if !mh.running {
+		return
+	}
+	var connectedVal float64 = 0
+	if connected {
+		connectedVal = 1
+	}
+	metric := mh.connected.MustCurryWith(prometheus.Labels{"jid": jid})
+	metric.With(prometheus.Labels{"bridge_connected": "true"}).Set(connectedVal)
+	metric.With(prometheus.Labels{"bridge_connected": "false"}).Set(1 - connectedVal)
+}
+
 func (mh *MetricsHandler) updateStats() {
 	start := time.Now()
 	var puppetCount int
diff --git a/provisioning.go b/provisioning.go
index 99d45cd..34fffce 100644
--- a/provisioning.go
+++ b/provisioning.go
@@ -111,6 +111,7 @@ func (prov *ProvisioningAPI) DeleteConnection(w http.ResponseWriter, r *http.Req
 	}
 	user.Conn.RemoveHandlers()
 	user.Conn = nil
+	user.bridge.Metrics.TrackConnectionState(user.JID, false)
 	jsonResponse(w, http.StatusOK, Response{true, "Disconnected from WhatsApp and connection deleted"})
 }
 
@@ -140,6 +141,7 @@ func (prov *ProvisioningAPI) Disconnect(w http.ResponseWriter, r *http.Request)
 	} else if len(sess.Wid) > 0 {
 		user.SetSession(&sess)
 	}
+	user.bridge.Metrics.TrackConnectionState(user.JID, false)
 	jsonResponse(w, http.StatusOK, Response{true, "Disconnected from WhatsApp"})
 }
 
diff --git a/user.go b/user.go
index 35e8a88..7fd169a 100644
--- a/user.go
+++ b/user.go
@@ -109,6 +109,7 @@ func (user *User) removeFromJIDMap() {
 	user.bridge.usersLock.Lock()
 	delete(user.bridge.usersByJID, user.JID)
 	user.bridge.usersLock.Unlock()
+	user.bridge.Metrics.TrackLoginState(user.JID, false)
 }
 
 func (bridge *Bridge) GetAllUsers() []*User {
@@ -407,6 +408,8 @@ func (cl ChatList) Swap(i, j int) {
 }
 
 func (user *User) PostLogin() {
+	user.bridge.Metrics.TrackConnectionState(user.JID, true)
+	user.bridge.Metrics.TrackLoginState(user.JID, true)
 	user.log.Debugln("Locking processing of incoming messages and starting post-login sync")
 	user.syncWait.Add(1)
 	user.syncStart <- struct{}{}
@@ -689,6 +692,7 @@ func (user *User) HandleError(err error) {
 	if closed, ok := err.(*whatsapp.ErrConnectionClosed); ok {
 		user.bridge.Metrics.TrackDisconnection(user.MXID)
 		if closed.Code == 1000 && user.cleanDisconnection {
+			user.bridge.Metrics.TrackConnectionState(user.JID, false)
 			user.cleanDisconnection = false
 			user.log.Infoln("Clean disconnection by server")
 			return
@@ -703,6 +707,7 @@ func (user *User) HandleError(err error) {
 }
 
 func (user *User) tryReconnect(msg string) {
+	user.bridge.Metrics.TrackConnectionState(user.JID, false)
 	if user.ConnectionErrors > user.bridge.Config.Bridge.MaxConnectionAttempts {
 		user.sendMarkdownBridgeAlert("%s. Use the `reconnect` command to reconnect.", msg)
 		return