chore: improve connection statistics (#49)

kevwan · Jul 11, 2023 · 9a3b558 · 9a3b558
1 parent 1f490c7
commit 9a3b558
Show file tree

Hide file tree

Showing 4 changed files with 28 additions and 13 deletions.
diff --git a/conn.go b/conn.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"io"
 	"net"
+	"strconv"
 	"sync"
 	"time"
 
@@ -81,7 +82,7 @@ func (c *PairedConnection) process() {
 
 	display.PrintlnWithTime(color.HiGreenString("[%d] Connected to server: %s", c.id, conn.RemoteAddr()))
 
-	stat.AddConn(fmt.Sprintf("%d:server", c.id), conn.(*net.TCPConn))
+	stat.AddConn(strconv.Itoa(c.id), conn.(*net.TCPConn))
 	c.svrConn = conn
 	go c.handleServerMessage()
 
@@ -91,8 +92,7 @@ func (c *PairedConnection) process() {
 func (c *PairedConnection) stop() {
 	c.once.Do(func() {
 		close(c.stopChan)
-		stat.DelConn(fmt.Sprintf("%d:server", c.id))
-		stat.DelConn(fmt.Sprintf("%d:client", c.id))
+		stat.DelConn(strconv.Itoa(c.id))
 
 		if c.cliConn != nil {
 			display.PrintlnWithTime(color.HiBlueString("[%d] Client connection closed", c.id))
@@ -128,7 +128,6 @@ func startListener() error {
 		display.PrintlnWithTime(color.HiGreenString("[%d] Accepted from: %s",
 			connIndex, cliConn.RemoteAddr()))
 
-		stat.AddConn(fmt.Sprintf("%d:client", connIndex), cliConn.(*net.TCPConn))
 		pconn := NewPairedConnection(connIndex, cliConn)
 		go pconn.process()
 	}

diff --git a/counter.go b/counter.go
@@ -25,7 +25,7 @@ func NewConnCounter() Stater {
 	}
 }
 
-func (c *connCounter) AddConn(key string, conn *net.TCPConn) {
+func (c *connCounter) AddConn(key string, _ *net.TCPConn) {
 	atomic.AddInt64(&c.total, 1)
 	val := atomic.AddInt64(&c.concurrent, 1)
 	max := atomic.LoadInt64(&c.max)
@@ -57,8 +57,18 @@ func (c *connCounter) Start() {
 }
 
 func (c *connCounter) Stop() {
+	c.lock.Lock()
+	for _, start := range c.conns {
+		lifetime := time.Since(start)
+		if lifetime > c.maxLifetime {
+			c.maxLifetime = lifetime
+		}
+	}
+	defer c.lock.Unlock()
+
 	fmt.Println()
-	color.HiWhite("Total connections: %d", atomic.LoadInt64(&c.total))
-	color.HiWhite("Max concurrent connections: %d", atomic.LoadInt64(&c.max))
-	color.HiWhite("Max connection lifetime: %s", c.maxLifetime)
+	color.HiWhite("Connection stats (client -> tproxy -> server):")
+	color.HiWhite("  Total connections: %d", atomic.LoadInt64(&c.total))
+	color.HiWhite("  Max concurrent connections: %d", atomic.LoadInt64(&c.max))
+	color.HiWhite("  Max connection lifetime: %s", c.maxLifetime)
 }
diff --git a/go.mod b/go.mod
@@ -12,8 +12,9 @@ require (
 
 require (
 	github.com/mattn/go-colorable v0.1.13 // indirect
-	github.com/mattn/go-isatty v0.0.17 // indirect
-	github.com/mattn/go-runewidth v0.0.9 // indirect
+	github.com/mattn/go-isatty v0.0.19 // indirect
+	github.com/mattn/go-runewidth v0.0.14 // indirect
+	github.com/rivo/uniseg v0.4.4 // indirect
 	golang.org/x/sys v0.10.0 // indirect
 	golang.org/x/text v0.11.0 // indirect
 )
diff --git a/go.sum b/go.sum
@@ -11,13 +11,17 @@ github.com/klauspost/compress v1.13.6/go.mod h1:/3/Vjq9QcHkK5uEr5lBEmyoZ1iFhe47e
 github.com/mattn/go-colorable v0.1.13 h1:fFA4WZxdEF4tXPZVKMLwD8oUnCTTo08duU7wxecdEvA=
 github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovkB8vQcUbaXHg=
 github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM=
-github.com/mattn/go-isatty v0.0.17 h1:BTarxUcIeDqL27Mc+vyvdWYSL28zpIhv3RoTdsLMPng=
-github.com/mattn/go-isatty v0.0.17/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM=
-github.com/mattn/go-runewidth v0.0.9 h1:Lm995f3rfxdpd6TSmuVCHVb/QhupuXlYr8sCI/QdE+0=
+github.com/mattn/go-isatty v0.0.19 h1:JITubQf0MOLdlGRuRq+jtsDlekdYPia9ZFsB8h/APPA=
+github.com/mattn/go-isatty v0.0.19/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
 github.com/mattn/go-runewidth v0.0.9/go.mod h1:H031xJmbD/WCDINGzjvQ9THkh0rPKHF+m2gUSrubnMI=
+github.com/mattn/go-runewidth v0.0.14 h1:+xnbZSEeDbOIg5/mE6JF0w6n9duR1l3/WmbinWVwUuU=
+github.com/mattn/go-runewidth v0.0.14/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w=
 github.com/montanaflynn/stats v0.0.0-20171201202039-1bf9dbcd8cbe/go.mod h1:wL8QJuTMNUDYhXwkmfOly8iTdp5TEcJFWZD2D7SIkUc=
 github.com/olekukonko/tablewriter v0.0.5 h1:P2Ga83D34wi1o9J6Wh1mRuqd4mF/x/lgBS7N7AbDhec=
 github.com/olekukonko/tablewriter v0.0.5/go.mod h1:hPp6KlRPjbx+hW8ykQs1w3UBbZlj6HuIJcUGPhkA7kY=
+github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc=
+github.com/rivo/uniseg v0.4.4 h1:8TfxU8dW6PdqD27gjM8MVNuicgxIjxpm4K7x4jp8sis=
+github.com/rivo/uniseg v0.4.4/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88=
 github.com/xdg-go/pbkdf2 v1.0.0/go.mod h1:jrpuAogTd400dnrH08LKmI/xc1MbPOebTwRqcT5RDeI=
 github.com/xdg-go/scram v1.1.2/go.mod h1:RT/sEzTbU5y00aCK8UOx6R7YryM0iF1N2MOmC3kKLN4=
 github.com/xdg-go/stringprep v1.0.4/go.mod h1:mPGuuIYwz7CmR2bT9j4GbQqutWS1zV24gijq1dTyGkM=
@@ -44,6 +48,7 @@ golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBc
 golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.10.0 h1:SqMFp9UcQJZa+pmYuAKjd9xq1f0j5rLcDIk0mj4qAsA=
 golang.org/x/sys v0.10.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=