From 5a5c93a756d282b2c34afd3cbbc2ba34bfb98cc6 Mon Sep 17 00:00:00 2001
From: Simon Ser <contact@emersion.fr>
Date: Mon, 15 Nov 2021 21:11:23 +0100
Subject: [PATCH] Add message counter metrics

---
 downstream.go | 10 ++++++++++
 server.go     | 25 +++++++++++++++++++++++++
 upstream.go   | 10 ++++++++++
 3 files changed, 45 insertions(+)

diff --git a/downstream.go b/downstream.go
index f57fa7e..d615f6a 100644
--- a/downstream.go
+++ b/downstream.go
@@ -444,6 +444,15 @@ func (dc *downstreamConn) unmarshalText(uc *upstreamConn, text string) string {
 	return strings.ReplaceAll(text, "/"+uc.network.GetName(), "")
 }
 
+func (dc *downstreamConn) ReadMessage() (*irc.Message, error) {
+	msg, err := dc.conn.ReadMessage()
+	if err != nil {
+		return nil, err
+	}
+	dc.srv.metrics.downstreamInMessagesTotal.Inc()
+	return msg, nil
+}
+
 func (dc *downstreamConn) readMessages(ch chan<- event) error {
 	for {
 		msg, err := dc.ReadMessage()
@@ -498,6 +507,7 @@ func (dc *downstreamConn) SendMessage(msg *irc.Message) {
 		return
 	}
 
+	dc.srv.metrics.downstreamOutMessagesTotal.Inc()
 	dc.conn.SendMessage(msg)
 }
 
diff --git a/server.go b/server.go
index 1fa7bb3..f77269d 100644
--- a/server.go
+++ b/server.go
@@ -101,6 +101,11 @@ type Server struct {
 	metrics struct {
 		downstreams int64Gauge
 		upstreams   int64Gauge
+
+		upstreamOutMessagesTotal   prometheus.Counter
+		upstreamInMessagesTotal    prometheus.Counter
+		downstreamOutMessagesTotal prometheus.Counter
+		downstreamInMessagesTotal  prometheus.Counter
 	}
 }
 
@@ -170,6 +175,26 @@ func (s *Server) registerMetrics() {
 		Name: "soju_upstreams_active",
 		Help: "Current number of upstream connections",
 	}, s.metrics.upstreams.Float64)
+
+	s.metrics.upstreamOutMessagesTotal = factory.NewCounter(prometheus.CounterOpts{
+		Name: "soju_upstream_out_messages_total",
+		Help: "Total number of outgoing messages sent to upstream servers",
+	})
+
+	s.metrics.upstreamInMessagesTotal = factory.NewCounter(prometheus.CounterOpts{
+		Name: "soju_upstream_in_messages_total",
+		Help: "Total number of incoming messages received from upstream servers",
+	})
+
+	s.metrics.downstreamOutMessagesTotal = factory.NewCounter(prometheus.CounterOpts{
+		Name: "soju_downstream_out_messages_total",
+		Help: "Total number of outgoing messages sent to downstream clients",
+	})
+
+	s.metrics.downstreamInMessagesTotal = factory.NewCounter(prometheus.CounterOpts{
+		Name: "soju_downstream_in_messages_total",
+		Help: "Total number of incoming messages received from downstream clients",
+	})
 }
 
 func (s *Server) Shutdown() {
diff --git a/upstream.go b/upstream.go
index 077d792..2f28412 100644
--- a/upstream.go
+++ b/upstream.go
@@ -1806,6 +1806,15 @@ func (uc *upstreamConn) register() {
 	})
 }
 
+func (uc *upstreamConn) ReadMessage() (*irc.Message, error) {
+	msg, err := uc.conn.ReadMessage()
+	if err != nil {
+		return nil, err
+	}
+	uc.srv.metrics.upstreamInMessagesTotal.Inc()
+	return msg, nil
+}
+
 func (uc *upstreamConn) runUntilRegistered() error {
 	for !uc.registered {
 		msg, err := uc.ReadMessage()
@@ -1856,6 +1865,7 @@ func (uc *upstreamConn) SendMessage(msg *irc.Message) {
 		msg.Tags = nil
 	}
 
+	uc.srv.metrics.upstreamOutMessagesTotal.Inc()
 	uc.conn.SendMessage(msg)
 }