From 2714ed97e7ca400d61053c46cf69a992d09c748a Mon Sep 17 00:00:00 2001
From: Huw Jones <huw@pexip.com>
Date: Thu, 4 Jul 2024 16:25:54 +0100
Subject: [PATCH] process_collector: collect in/out bytes

Signed-off-by: Huw Jones <huw@pexip.com>
---
 prometheus/process_collector.go       | 27 +++++++++++++++++++--------
 prometheus/process_collector_other.go | 14 ++++++++++++++
 prometheus/process_collector_test.go  |  4 ++++
 3 files changed, 37 insertions(+), 8 deletions(-)

diff --git a/prometheus/process_collector.go b/prometheus/process_collector.go
index 8548dd18e..e217f262d 100644
--- a/prometheus/process_collector.go
+++ b/prometheus/process_collector.go
@@ -22,14 +22,15 @@ import (
 )
 
 type processCollector struct {
-	collectFn       func(chan<- Metric)
-	pidFn           func() (int, error)
-	reportErrors    bool
-	cpuTotal        *Desc
-	openFDs, maxFDs *Desc
-	vsize, maxVsize *Desc
-	rss             *Desc
-	startTime       *Desc
+	collectFn         func(chan<- Metric)
+	pidFn             func() (int, error)
+	reportErrors      bool
+	cpuTotal          *Desc
+	openFDs, maxFDs   *Desc
+	vsize, maxVsize   *Desc
+	rss               *Desc
+	startTime         *Desc
+	inBytes, outBytes *Desc
 }
 
 // ProcessCollectorOpts defines the behavior of a process metrics collector
@@ -100,6 +101,16 @@ func NewProcessCollector(opts ProcessCollectorOpts) Collector {
 			"Start time of the process since unix epoch in seconds.",
 			nil, nil,
 		),
+		inBytes: NewDesc(
+			ns+"process_network_bytes_in_total",
+			"Number of bytes received by the process over the network.",
+			nil, nil,
+		),
+		outBytes: NewDesc(
+			ns+"process_network_bytes_out_total",
+			"Number of bytes sent by the process over the network.",
+			nil, nil,
+		),
 	}
 
 	if opts.PidFn == nil {
diff --git a/prometheus/process_collector_other.go b/prometheus/process_collector_other.go
index 8c1136cee..14d56d2d0 100644
--- a/prometheus/process_collector_other.go
+++ b/prometheus/process_collector_other.go
@@ -63,4 +63,18 @@ func (c *processCollector) processCollect(ch chan<- Metric) {
 	} else {
 		c.reportError(ch, nil, err)
 	}
+
+	if netstat, err := p.Netstat(); err == nil {
+		var inOctets, outOctets float64
+		if netstat.IpExt.InOctets != nil {
+			inOctets = *netstat.IpExt.InOctets
+		}
+		if netstat.IpExt.OutOctets != nil {
+			outOctets = *netstat.IpExt.OutOctets
+		}
+		ch <- MustNewConstMetric(c.inBytes, CounterValue, inOctets)
+		ch <- MustNewConstMetric(c.outBytes, CounterValue, outOctets)
+	} else {
+		c.reportError(ch, nil, err)
+	}
 }
diff --git a/prometheus/process_collector_test.go b/prometheus/process_collector_test.go
index 3a604aba9..025951725 100644
--- a/prometheus/process_collector_test.go
+++ b/prometheus/process_collector_test.go
@@ -69,6 +69,8 @@ func TestProcessCollector(t *testing.T) {
 		regexp.MustCompile("\nprocess_virtual_memory_bytes [1-9]"),
 		regexp.MustCompile("\nprocess_resident_memory_bytes [1-9]"),
 		regexp.MustCompile("\nprocess_start_time_seconds [0-9.]{10,}"),
+		regexp.MustCompile("\nprocess_net_bytes_in [0-9]+"),
+		regexp.MustCompile("\nprocess_net_bytes_out [0-9]+"),
 		regexp.MustCompile("\nfoobar_process_cpu_seconds_total [0-9]"),
 		regexp.MustCompile("\nfoobar_process_max_fds [1-9]"),
 		regexp.MustCompile("\nfoobar_process_open_fds [1-9]"),
@@ -76,6 +78,8 @@ func TestProcessCollector(t *testing.T) {
 		regexp.MustCompile("\nfoobar_process_virtual_memory_bytes [1-9]"),
 		regexp.MustCompile("\nfoobar_process_resident_memory_bytes [1-9]"),
 		regexp.MustCompile("\nfoobar_process_start_time_seconds [0-9.]{10,}"),
+		regexp.MustCompile("\nfoobar_process_net_bytes_in [0-9]+"),
+		regexp.MustCompile("\nfoobar_process_net_bytes_out [0-9]+"),
 	} {
 		if !re.Match(buf.Bytes()) {
 			t.Errorf("want body to match %s\n%s", re, buf.String())