From e60b09eb518a89e79a7938f5b4891b7fd08790a9 Mon Sep 17 00:00:00 2001
From: Elle Mouton <elle.mouton@gmail.com>
Date: Thu, 31 Mar 2022 11:27:10 +0200
Subject: [PATCH] multi: add mailbox read counter

In this commit, we add a mailbox-read-count metric. This will be
incremented each time a mailbox with an _odd_ stream ID is read from.
We do this because we assume that a full duplex connection is being used
meaning that there will be 2 streams that have a matching ID except for
the last byte. And so to avoid duplicating the data, we only record the
odd streams. We also assume that for every read, there will be a write
and so we only record the reads.
---
 hashmail_server.go | 27 +++++++++++++++++++++++++++
 prometheus.go      | 17 +++++++++++++++++
 2 files changed, 44 insertions(+)

diff --git a/hashmail_server.go b/hashmail_server.go
index 9378726..cb432d4 100644
--- a/hashmail_server.go
+++ b/hashmail_server.go
@@ -11,6 +11,7 @@ import (
 
 	"github.com/lightninglabs/lightning-node-connect/hashmailrpc"
 	"github.com/lightningnetwork/lnd/tlv"
+	"github.com/prometheus/client_golang/prometheus"
 	"golang.org/x/time/rate"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -43,6 +44,19 @@ func newStreamID(id []byte) streamID {
 	return s
 }
 
+// baseID returns the first 16 bytes of the streamID. This part of the ID will
+// overlap for the two streams in a bidirectional pair.
+func (s *streamID) baseID() [16]byte {
+	var id [16]byte
+	copy(id[:], s[:16])
+	return id
+}
+
+// isOdd returns true if the streamID is an odd number.
+func (s *streamID) isOdd() bool {
+	return s[63]&0x01 == 0x01
+}
+
 // readStream is the read side of the read pipe, which is implemented a
 // buffered wrapper around the core reader.
 type readStream struct {
@@ -670,6 +684,19 @@ func (h *hashMailServer) RecvStream(desc *hashmailrpc.CipherBoxDesc,
 		log.Tracef("Read %v bytes for HashMail stream_id=%x",
 			len(nextMsg), desc.StreamId)
 
+		// In order not to duplicate metric data, we only record this
+		// read if its streamID is odd. We use the base stream ID as the
+		// label. For this to work, it is expected that the read and
+		// write streams of bidirectional pair have the same IDs with
+		// the last bit flipped for one of them.
+		streamID := newStreamID(desc.StreamId)
+		if streamID.isOdd() {
+			baseID := streamID.baseID()
+			mailboxReadCount.With(prometheus.Labels{
+				streamIDLabel: fmt.Sprintf("%x", baseID),
+			}).Inc()
+		}
+
 		err = reader.Send(&hashmailrpc.CipherBox{
 			Desc: desc,
 			Msg:  nextMsg,
diff --git a/prometheus.go b/prometheus.go
index 99b8eff..1585b3f 100644
--- a/prometheus.go
+++ b/prometheus.go
@@ -8,12 +8,28 @@ import (
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 )
 
+const streamIDLabel = "streamID"
+
 var (
 	// mailboxCount tracks the current number of active mailboxes.
 	mailboxCount = prometheus.NewGauge(prometheus.GaugeOpts{
 		Namespace: "hashmail",
 		Name:      "mailbox_count",
 	})
+
+	// mailboxReadCount counts each time a mailbox pair is being used.
+	// A session consists of a bidirectional stream each using a mailbox
+	// with an ID that overlaps for the first 63 bytes and differ for the
+	// last bit. So in order to obtain accurate data about a specific
+	// mailbox session, the stream ID that will be recorded is the first
+	// 16 bytes of the session ID and we will only record the odd stream's
+	// reads so that we don't duplicate the data.
+	mailboxReadCount = prometheus.NewCounterVec(
+		prometheus.CounterOpts{
+			Namespace: "hashmail",
+			Name:      "mailbox_read_count",
+		}, []string{streamIDLabel},
+	)
 )
 
 // PrometheusConfig is the set of configuration data that specifies if
@@ -39,6 +55,7 @@ func StartPrometheusExporter(cfg *PrometheusConfig) error {
 
 	// Next, we'll register all our metrics.
 	prometheus.MustRegister(mailboxCount)
+	prometheus.MustRegister(mailboxReadCount)
 
 	// Finally, we'll launch the HTTP server that Prometheus will use to
 	// scape our metrics.