From 5a04f6c87182412d26188f545f728faffeb1a286 Mon Sep 17 00:00:00 2001
From: Tulir Asokan <tulir@maunium.net>
Date: Wed, 17 Jun 2020 17:57:14 +0300
Subject: [PATCH] Add simple disconnection count to metrics

---
 metrics.go | 18 ++++++++++++++++++
 user.go    |  2 ++
 2 files changed, 20 insertions(+)

diff --git a/metrics.go b/metrics.go
index 2751266..35a566b 100644
--- a/metrics.go
+++ b/metrics.go
@@ -27,6 +27,7 @@ import (
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 	log "maunium.net/go/maulogger/v2"
 	"maunium.net/go/mautrix/event"
+	"maunium.net/go/mautrix/id"
 
 	"maunium.net/go/mautrix-whatsapp/database"
 )
@@ -42,6 +43,7 @@ type MetricsHandler struct {
 
 	messageHandling         *prometheus.HistogramVec
 	countCollection         prometheus.Histogram
+	disconnections          *prometheus.CounterVec
 	puppetCount             prometheus.Gauge
 	userCount               prometheus.Gauge
 	messageCount            prometheus.Gauge
@@ -71,6 +73,10 @@ func NewMetricsHandler(address string, log log.Logger, db *database.Database) *M
 			Name: "whatsapp_count_collection",
 			Help: "Time spent collecting the whatsapp_*_total metrics",
 		}),
+		disconnections: promauto.NewCounterVec(prometheus.CounterOpts{
+			Name: "whatsapp_disconnections",
+			Help: "Number of times a Matrix user has been disconnected from WhatsApp",
+		}, []string{"user_id"}),
 		puppetCount: promauto.NewGauge(prometheus.GaugeOpts{
 			Name: "whatsapp_puppets_total",
 			Help: "Number of WhatsApp users bridged into Matrix",
@@ -91,7 +97,12 @@ func NewMetricsHandler(address string, log log.Logger, db *database.Database) *M
 	}
 }
 
+func noop() {}
+
 func (mh *MetricsHandler) TrackEvent(eventType event.Type) func() {
+	if !mh.running {
+		return noop
+	}
 	start := time.Now()
 	return func() {
 		duration := time.Now().Sub(start)
@@ -101,6 +112,13 @@ func (mh *MetricsHandler) TrackEvent(eventType event.Type) func() {
 	}
 }
 
+func (mh *MetricsHandler) TrackDisconnection(userID id.UserID) {
+	if !mh.running {
+		return
+	}
+	mh.disconnections.With(prometheus.Labels{"user_id": string(userID)}).Inc()
+}
+
 func (mh *MetricsHandler) updateStats() {
 	start := time.Now()
 	var puppetCount int
diff --git a/user.go b/user.go
index 3bda231..14eda22 100644
--- a/user.go
+++ b/user.go
@@ -561,6 +561,7 @@ func (user *User) HandleError(err error) {
 		user.log.Errorfln("WhatsApp error: %v", err)
 	}
 	if closed, ok := err.(*whatsapp.ErrConnectionClosed); ok {
+		user.bridge.Metrics.TrackDisconnection(user.MXID)
 		if closed.Code == 1000 && user.cleanDisconnection {
 			user.cleanDisconnection = false
 			user.log.Infoln("Clean disconnection by server")
@@ -568,6 +569,7 @@ func (user *User) HandleError(err error) {
 		}
 		go user.tryReconnect(fmt.Sprintf("Your WhatsApp connection was closed with websocket status code %d", closed.Code))
 	} else if failed, ok := err.(*whatsapp.ErrConnectionFailed); ok {
+		user.bridge.Metrics.TrackDisconnection(user.MXID)
 		user.ConnectionErrors++
 		go user.tryReconnect(fmt.Sprintf("Your WhatsApp connection failed: %v", failed.Err))
 	}