feat: support connection stats (#36)

conn.go

@@ -106,7 +106,7 @@ func (c *PairedConnection) stop() {
 }
 
 func startListener() error {
-	stat = NewStatPrinter(statInterval)
+	stat = NewStater(NewConnCounter(), NewStatPrinter(statInterval))
 	go stat.Start()
 
 	conn, err := net.Listen("tcp", fmt.Sprintf("%s:%d", settings.LocalHost, settings.LocalPort))

counter.go

@@ -0,0 +1,64 @@
+package main
+
+import (
+	"fmt"
+	"net"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	"github.com/fatih/color"
+)
+
+type connCounter struct {
+	total       int64
+	concurrent  int64
+	max         int64
+	conns       map[string]time.Time
+	maxLifetime time.Duration
+	lock        sync.Mutex
+}
+
+func NewConnCounter() Stater {
+	return &connCounter{
+		conns: make(map[string]time.Time),
+	}
+}
+
+func (c *connCounter) AddConn(key string, conn *net.TCPConn) {
+	atomic.AddInt64(&c.total, 1)
+	val := atomic.AddInt64(&c.concurrent, 1)
+	max := atomic.LoadInt64(&c.max)
+	if val > max {
+		atomic.CompareAndSwapInt64(&c.max, max, val)
+	}
+
+	c.lock.Lock()
+	defer c.lock.Unlock()
+	c.conns[key] = time.Now()
+}
+
+func (c *connCounter) DelConn(key string) {
+	atomic.AddInt64(&c.concurrent, -1)
+
+	c.lock.Lock()
+	defer c.lock.Unlock()
+	start, ok := c.conns[key]
+	delete(c.conns, key)
+	if ok {
+		lifetime := time.Since(start)
+		if lifetime > c.maxLifetime {
+			c.maxLifetime = lifetime
+		}
+	}
+}
+
+func (c *connCounter) Start() {
+}
+
+func (c *connCounter) Stop() {
+	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)
+}

stat+polyfill.go

@@ -21,3 +21,6 @@ func (p StatPrinter) DelConn(_ string) {
 
 func (p StatPrinter) Start() {
 }
+
+func (p StatPrinter) Stop() {
+}

stat.go

@@ -1,14 +1,80 @@
 package main
 
 import (
+	"fmt"
 	"net"
+	"os"
+	"os/signal"
+	"syscall"
 	"time"
 )
 
-type Stater interface {
-	AddConn(key string, conn *net.TCPConn)
-	DelConn(key string)
-	Start()
+type (
+	Stater interface {
+		AddConn(key string, conn *net.TCPConn)
+		DelConn(key string)
+		Start()
+		Stop()
+	}
+
+	compositeStater struct {
+		staters []Stater
+	}
+)
+
+func NewStater(staters ...Stater) Stater {
+	stat := compositeStater{
+		staters: append([]Stater(nil), staters...),
+	}
+
+	go func() {
+		c := make(chan os.Signal, 1)
+		signal.Notify(c, os.Interrupt, syscall.SIGTERM)
+
+		for {
+			select {
+			case sig := <-c:
+				signal.Stop(c)
+				stat.Stop()
+
+				p, err := os.FindProcess(syscall.Getpid())
+				if err != nil {
+					fmt.Println(err)
+					os.Exit(0)
+				}
+
+				if err := p.Signal(sig); err != nil {
+					fmt.Println(err)
+				}
+			}
+		}
+	}()
+
+	return stat
+}
+
+func (c compositeStater) AddConn(key string, conn *net.TCPConn) {
+	for _, s := range c.staters {
+		s.AddConn(key, conn)
+	}
+}
+
+func (c compositeStater) DelConn(key string) {
+	for _, s := range c.staters {
+		s.DelConn(key)
+	}
+}
+
+func (c compositeStater) Start() {
+	for _, s := range c.staters {
+		s.Start()
+	}
+}
+
+func (c compositeStater) Stop() {
+	for _, s := range c.staters {
+		s.Stop()
+	}
 }
 
 type NilPrinter struct{}
@@ -25,3 +91,6 @@ func (p NilPrinter) DelConn(_ string) {
 
 func (p NilPrinter) Start() {
 }
+
+func (p NilPrinter) Stop() {
+}

stat_linux.go

@@ -51,6 +51,10 @@ func (p *StatPrinter) Start() {
 	}
 }
 
+func (p *StatPrinter) Stop() {
+	p.print()
+}
+
 func (p *StatPrinter) buildRows() [][]string {
 	var keys []string
 	infos := make(map[string]*TcpInfo)