From d232ccbfab351c93341ff5eeda5413394d6da938 Mon Sep 17 00:00:00 2001
From: Wade Simmons <wadey@slack-corp.com>
Date: Mon, 1 Mar 2021 19:51:33 -0500
Subject: [PATCH] add metrics for the udp sockets using SO_MEMINFO  (#390)

Retrieve the current socket stats using SO_MEMINFO and report them as
metrics gauges. If SO_MEMINFO isn't supported, we don't report these metrics.
---
 interface.go   |  5 +++++
 udp_generic.go |  5 +++++
 udp_linux.go   | 59 ++++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 69 insertions(+)

diff --git a/interface.go b/interface.go
index 8422187..825ba97 100644
--- a/interface.go
+++ b/interface.go
@@ -268,8 +268,13 @@ func (f *Interface) reloadFirewall(c *Config) {
 
 func (f *Interface) emitStats(i time.Duration) {
 	ticker := time.NewTicker(i)
+
+	udpStats := NewUDPStatsEmitter(f.writers)
+
 	for range ticker.C {
 		f.firewall.EmitStats()
 		f.handshakeManager.EmitStats()
+
+		udpStats()
 	}
 }
diff --git a/udp_generic.go b/udp_generic.go
index eccf616..5a1d204 100644
--- a/udp_generic.go
+++ b/udp_generic.go
@@ -96,6 +96,11 @@ func (u *udpConn) reloadConfig(c *Config) {
 	// TODO
 }
 
+func NewUDPStatsEmitter(udpConns []*udpConn) func() {
+	// No UDP stats for non-linux
+	return func() {}
+}
+
 type rawMessage struct {
 	Len uint32
 }
diff --git a/udp_linux.go b/udp_linux.go
index 19b5f5e..69eee31 100644
--- a/udp_linux.go
+++ b/udp_linux.go
@@ -12,6 +12,7 @@ import (
 	"syscall"
 	"unsafe"
 
+	"github.com/rcrowley/go-metrics"
 	"golang.org/x/sys/unix"
 )
 
@@ -55,6 +56,23 @@ type rawSockaddrAny struct {
 
 var x int
 
+// From linux/sock_diag.h
+const (
+	_SK_MEMINFO_RMEM_ALLOC = iota
+	_SK_MEMINFO_RCVBUF
+	_SK_MEMINFO_WMEM_ALLOC
+	_SK_MEMINFO_SNDBUF
+	_SK_MEMINFO_FWD_ALLOC
+	_SK_MEMINFO_WMEM_QUEUED
+	_SK_MEMINFO_OPTMEM
+	_SK_MEMINFO_BACKLOG
+	_SK_MEMINFO_DROPS
+
+	_SK_MEMINFO_VARS
+)
+
+type _SK_MEMINFO [_SK_MEMINFO_VARS]uint32
+
 func NewListener(ip string, port int, multi bool) (*udpConn, error) {
 	syscall.ForkLock.RLock()
 	fd, err := unix.Socket(unix.AF_INET, unix.SOCK_DGRAM, unix.IPPROTO_UDP)
@@ -281,6 +299,47 @@ func (u *udpConn) reloadConfig(c *Config) {
 	}
 }
 
+func (u *udpConn) getMemInfo(meminfo *_SK_MEMINFO) error {
+	var vallen uint32 = 4 * _SK_MEMINFO_VARS
+	_, _, err := unix.Syscall6(unix.SYS_GETSOCKOPT, uintptr(u.sysFd), uintptr(unix.SOL_SOCKET), uintptr(unix.SO_MEMINFO), uintptr(unsafe.Pointer(meminfo)), uintptr(unsafe.Pointer(&vallen)), 0)
+	if err != 0 {
+		return err
+	}
+	return nil
+}
+
+func NewUDPStatsEmitter(udpConns []*udpConn) func() {
+	// Check if our kernel supports SO_MEMINFO before registering the gauges
+	var udpGauges [][_SK_MEMINFO_VARS]metrics.Gauge
+	var meminfo _SK_MEMINFO
+	if err := udpConns[0].getMemInfo(&meminfo); err == nil {
+		udpGauges = make([][_SK_MEMINFO_VARS]metrics.Gauge, len(udpConns))
+		for i := range udpConns {
+			udpGauges[i] = [_SK_MEMINFO_VARS]metrics.Gauge{
+				metrics.GetOrRegisterGauge(fmt.Sprintf("udp.%d.rmem_alloc", i), nil),
+				metrics.GetOrRegisterGauge(fmt.Sprintf("udp.%d.rcvbuf", i), nil),
+				metrics.GetOrRegisterGauge(fmt.Sprintf("udp.%d.wmem_alloc", i), nil),
+				metrics.GetOrRegisterGauge(fmt.Sprintf("udp.%d.sndbuf", i), nil),
+				metrics.GetOrRegisterGauge(fmt.Sprintf("udp.%d.fwd_alloc", i), nil),
+				metrics.GetOrRegisterGauge(fmt.Sprintf("udp.%d.wmem_queued", i), nil),
+				metrics.GetOrRegisterGauge(fmt.Sprintf("udp.%d.optmem", i), nil),
+				metrics.GetOrRegisterGauge(fmt.Sprintf("udp.%d.backlog", i), nil),
+				metrics.GetOrRegisterGauge(fmt.Sprintf("udp.%d.drops", i), nil),
+			}
+		}
+	}
+
+	return func() {
+		for i, gauges := range udpGauges {
+			if err := udpConns[i].getMemInfo(&meminfo); err == nil {
+				for j := 0; j < _SK_MEMINFO_VARS; j++ {
+					gauges[j].Update(int64(meminfo[j]))
+				}
+			}
+		}
+	}
+}
+
 func (ua *udpAddr) Equals(t *udpAddr) bool {
 	if t == nil || ua == nil {
 		return t == nil && ua == nil