From 0cc52eee460c9ee9c5f495c1cebda4f8789e14b1 Mon Sep 17 00:00:00 2001
From: Shizun Ge <shizunge@gmail.com>
Date: Sat, 27 Jan 2024 23:26:39 -0800
Subject: [PATCH] Prometheus roundup the trapped time to the interval.

In the old way, if the connection is broken less than an interval, the trapped time won't be reported.
With this fix, the prometheus should report the same value as the log.
---
 main.go            | 11 ++++++-----
 metrics/metrics.go |  9 +++++++--
 2 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/main.go b/main.go
index 64b996e..c9d0353 100644
--- a/main.go
+++ b/main.go
@@ -42,23 +42,24 @@ func startSending(maxClients int64, bannerMaxLength int64, records chan<- metric
 				bytesSent, err := c.Send(bannerMaxLength)
 				remoteIpAddr := c.RemoteIpAddr()
 				localPort := c.LocalPort()
+				millisecondsSpent := c.MillisecondsSinceLast()
 				if err != nil {
 					c.Close()
 					records <- metrics.RecordEntry{
-						RecordType: metrics.RecordEntryTypeStop,
-						IpAddr:     remoteIpAddr,
-						LocalPort:  localPort,
+						RecordType:        metrics.RecordEntryTypeStop,
+						IpAddr:            remoteIpAddr,
+						LocalPort:         localPort,
+						MillisecondsSpent: millisecondsSpent,
 					}
 					return
 				}
-				millisecondsSpent := c.MillisecondsSinceLast()
 				clients <- c
 				records <- metrics.RecordEntry{
 					RecordType:        metrics.RecordEntryTypeSend,
 					IpAddr:            remoteIpAddr,
 					LocalPort:         localPort,
-					BytesSent:         bytesSent,
 					MillisecondsSpent: millisecondsSpent,
+					BytesSent:         bytesSent,
 				}
 			}()
 		}
diff --git a/metrics/metrics.go b/metrics/metrics.go
index fc6773b..c754360 100644
--- a/metrics/metrics.go
+++ b/metrics/metrics.go
@@ -106,8 +106,8 @@ type RecordEntry struct {
 	RecordType        int
 	IpAddr            string
 	LocalPort         string
-	BytesSent         int
 	MillisecondsSpent int64
+	BytesSent         int
 }
 
 func StartRecording(maxClients int64, prometheusEnabled bool, prometheusCleanUnseenSeconds int, geoOption geoip.GeoOption) chan RecordEntry {
@@ -140,10 +140,15 @@ func StartRecording(maxClients int64, prometheusEnabled bool, prometheusCleanUns
 				clientSeconds.With(prometheus.Labels{
 					"ip":         r.IpAddr,
 					"local_port": r.LocalPort}).Add(secondsSpent)
-				totalBytes.With(prometheus.Labels{"local_port": r.LocalPort}).Add(float64(r.BytesSent))
 				totalSeconds.With(prometheus.Labels{"local_port": r.LocalPort}).Add(secondsSpent)
+				totalBytes.With(prometheus.Labels{"local_port": r.LocalPort}).Add(float64(r.BytesSent))
 				pq.Update(r.IpAddr, time.Now())
 			case RecordEntryTypeStop:
+				secondsSpent := float64(r.MillisecondsSpent) / 1000
+				clientSeconds.With(prometheus.Labels{
+					"ip":         r.IpAddr,
+					"local_port": r.LocalPort}).Add(secondsSpent)
+				totalSeconds.With(prometheus.Labels{"local_port": r.LocalPort}).Add(secondsSpent)
 				totalClientsClosed.With(prometheus.Labels{"local_port": r.LocalPort}).Inc()
 				pq.Update(r.IpAddr, time.Now())
 			case RecordEntryTypeClean: