From 136f232c00e5e1b09dfb04b204b17bcfddf8396e Mon Sep 17 00:00:00 2001
From: Devin Carr <dcarr@cloudflare.com>
Date: Fri, 16 Jun 2023 17:07:56 -0700
Subject: [PATCH] TUN-7477: Add UDP/TCP session metrics

New gauge metrics are exposed in the prometheus endpoint to
capture the current and total TCP and UDP sessions that
cloudflared has proxied.
---
 datagramsession/manager.go |  2 ++
 datagramsession/metrics.go | 40 ++++++++++++++++++++++++++++++++++++++
 proxy/metrics.go           | 29 +++++++++++++++++++++++++++
 proxy/proxy.go             |  4 ++--
 4 files changed, 73 insertions(+), 2 deletions(-)
 create mode 100644 datagramsession/metrics.go

diff --git a/datagramsession/manager.go b/datagramsession/manager.go
index d1062f33..60b701ab 100644
--- a/datagramsession/manager.go
+++ b/datagramsession/manager.go
@@ -118,6 +118,7 @@ func (m *manager) registerSession(ctx context.Context, registration *registerSes
 	session := m.newSession(registration.sessionID, registration.originProxy)
 	m.sessions[registration.sessionID] = session
 	registration.resultChan <- session
+	incrementUDPSessions()
 }
 
 func (m *manager) newSession(id uuid.UUID, dstConn io.ReadWriteCloser) *Session {
@@ -163,6 +164,7 @@ func (m *manager) unregisterSession(unregistration *unregisterSessionEvent) {
 	if ok {
 		delete(m.sessions, unregistration.sessionID)
 		session.close(unregistration.err)
+		decrementUDPActiveSessions()
 	}
 }
 
diff --git a/datagramsession/metrics.go b/datagramsession/metrics.go
new file mode 100644
index 00000000..12844e7a
--- /dev/null
+++ b/datagramsession/metrics.go
@@ -0,0 +1,40 @@
+package datagramsession
+
+import (
+	"github.com/prometheus/client_golang/prometheus"
+)
+
+const (
+	namespace = "cloudflared"
+)
+
+var (
+	activeUDPSessions = prometheus.NewGauge(prometheus.GaugeOpts{
+		Namespace: namespace,
+		Subsystem: "udp",
+		Name:      "active_sessions",
+		Help:      "Concurrent count of UDP sessions that are being proxied to any origin",
+	})
+	totalUDPSessions = prometheus.NewGauge(prometheus.GaugeOpts{
+		Namespace: namespace,
+		Subsystem: "udp",
+		Name:      "total_sessions",
+		Help:      "Total count of UDP sessions that have been proxied to any origin",
+	})
+)
+
+func init() {
+	prometheus.MustRegister(
+		activeUDPSessions,
+		totalUDPSessions,
+	)
+}
+
+func incrementUDPSessions() {
+	totalUDPSessions.Inc()
+	activeUDPSessions.Inc()
+}
+
+func decrementUDPActiveSessions() {
+	activeUDPSessions.Dec()
+}
diff --git a/proxy/metrics.go b/proxy/metrics.go
index e5406681..33a51721 100644
--- a/proxy/metrics.go
+++ b/proxy/metrics.go
@@ -43,6 +43,22 @@ var (
 			Help:      "Count of error proxying to origin",
 		},
 	)
+	activeTCPSessions = prometheus.NewGauge(
+		prometheus.GaugeOpts{
+			Namespace: connection.MetricsNamespace,
+			Subsystem: "tcp",
+			Name:      "active_sessions",
+			Help:      "Concurrent count of TCP sessions that are being proxied to any origin",
+		},
+	)
+	totalTCPSessions = prometheus.NewCounter(
+		prometheus.CounterOpts{
+			Namespace: connection.MetricsNamespace,
+			Subsystem: "tcp",
+			Name:      "total_sessions",
+			Help:      "Total count of TCP sessions that have been proxied to any origin",
+		},
+	)
 )
 
 func init() {
@@ -51,6 +67,8 @@ func init() {
 		concurrentRequests,
 		responseByCode,
 		requestErrors,
+		activeTCPSessions,
+		totalTCPSessions,
 	)
 }
 
@@ -62,3 +80,14 @@ func incrementRequests() {
 func decrementConcurrentRequests() {
 	concurrentRequests.Dec()
 }
+
+func incrementTCPRequests() {
+	incrementRequests()
+	totalTCPSessions.Inc()
+	activeTCPSessions.Inc()
+}
+
+func decrementTCPConcurrentRequests() {
+	decrementConcurrentRequests()
+	activeTCPSessions.Dec()
+}
diff --git a/proxy/proxy.go b/proxy/proxy.go
index 1664ead3..18e53dc2 100644
--- a/proxy/proxy.go
+++ b/proxy/proxy.go
@@ -158,8 +158,8 @@ func (p *Proxy) ProxyTCP(
 	rwa connection.ReadWriteAcker,
 	req *connection.TCPRequest,
 ) error {
-	incrementRequests()
-	defer decrementConcurrentRequests()
+	incrementTCPRequests()
+	defer decrementTCPConcurrentRequests()
 
 	if p.warpRouting == nil {
 		err := errors.New(`cloudflared received a request from WARP client, but your configuration has disabled ingress from WARP clients. To enable this, set "warp-routing:\n\t enabled: true" in your config.yaml`)