From 830bec4a4209b359df165774a49ee2fa954e42a0 Mon Sep 17 00:00:00 2001
From: Erope <wjhjd163@gmail.com>
Date: Thu, 30 Dec 2021 18:01:32 +0800
Subject: [PATCH] Add SS Func to accelerate connections count in some Linux

---
 cmd/agent/monitor/monitor.go | 33 +++++++++++++++++++++++----------
 go.mod                       |  1 +
 go.sum                       |  2 ++
 3 files changed, 26 insertions(+), 10 deletions(-)

diff --git a/cmd/agent/monitor/monitor.go b/cmd/agent/monitor/monitor.go
index 855601b3e7..3e9e5822d3 100644
--- a/cmd/agent/monitor/monitor.go
+++ b/cmd/agent/monitor/monitor.go
@@ -2,14 +2,15 @@ package monitor
 
 import (
 	"fmt"
+	"os/exec"
 	"regexp"
 	"runtime"
+	"strconv"
 	"strings"
 	"syscall"
 	"time"
-	"strconv"
-	"os/exec"
 
+	"github.com/Erope/goss"
 	"github.com/shirou/gopsutil/v3/cpu"
 	"github.com/shirou/gopsutil/v3/disk"
 	"github.com/shirou/gopsutil/v3/host"
@@ -114,13 +115,25 @@ func GetState(skipConnectionCount bool, skipProcsCount bool) *model.HostState {
 
 	var tcpConnCount, udpConnCount uint64
 	if !skipConnectionCount {
-		conns, _ := net.Connections("all")
-		for i := 0; i < len(conns); i++ {
-			switch conns[i].Type {
-			case syscall.SOCK_STREAM:
-				tcpConnCount++
-			case syscall.SOCK_DGRAM:
-				udpConnCount++
+		ss_err := true
+		if runtime.GOOS == "linux" {
+			tcpStat, err_tcp := goss.ConnectionsWithProtocol(syscall.IPPROTO_TCP)
+			udpStat, err_udp := goss.ConnectionsWithProtocol(syscall.IPPROTO_UDP)
+			if err_tcp == nil && err_udp == nil {
+				ss_err = false
+				tcpConnCount = uint64(len(tcpStat))
+				udpConnCount = uint64(len(udpStat))
+			}
+		}
+		if ss_err {
+			conns, _ := net.Connections("all")
+			for i := 0; i < len(conns); i++ {
+				switch conns[i].Type {
+				case syscall.SOCK_STREAM:
+					tcpConnCount++
+				case syscall.SOCK_DGRAM:
+					udpConnCount++
+				}
 			}
 		}
 	}
@@ -192,7 +205,7 @@ func getDiskTotalAndUsed() (total uint64, used uint64) {
 		}
 		used += diskUsageOf.Used
 	}
-	
+
 	// Fallback 到这个方法,仅统计根路径,适用于OpenVZ之类的.
 	if runtime.GOOS == "linux" {
 		if total == 0 && used == 0 {
diff --git a/go.mod b/go.mod
index 175bea8537..1091f4c704 100644
--- a/go.mod
+++ b/go.mod
@@ -4,6 +4,7 @@ go 1.13
 
 require (
 	code.cloudfoundry.org/bytefmt v0.0.0-20211005130812-5bb3c17173e5
+	github.com/Erope/goss v0.0.0-20211230093305-df3c03fd1ed4 // indirect
 	github.com/artdarek/go-unzip v1.0.0
 	github.com/blang/semver v3.5.1+incompatible
 	github.com/creack/pty v1.1.17
diff --git a/go.sum b/go.sum
index f09c8c86ee..ea016dda9a 100644
--- a/go.sum
+++ b/go.sum
@@ -46,6 +46,8 @@ code.cloudfoundry.org/bytefmt v0.0.0-20211005130812-5bb3c17173e5/go.mod h1:v4VVB
 dmitri.shuralyov.com/gpu/mtl v0.0.0-20190408044501-666a987793e9/go.mod h1:H6x//7gZCb22OMCxBHrMx7a5I7Hp++hsVxbQ4BYO7hU=
 github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
 github.com/BurntSushi/xgb v0.0.0-20160522181843-27f122750802/go.mod h1:IVnqGOEym/WlBOVXweHU+Q+/VP0lqqI8lqeDx9IjBqo=
+github.com/Erope/goss v0.0.0-20211230093305-df3c03fd1ed4 h1:X35U3bryt+j2a9RZyXyBPISNIYQsfNKWC/d/B8J861I=
+github.com/Erope/goss v0.0.0-20211230093305-df3c03fd1ed4/go.mod h1:gl55GUYOV6rvsL/V23GdfRNdnreh+McD5Yo1sNc+Qe0=
 github.com/OneOfOne/xxhash v1.2.2/go.mod h1:HSdplMjZKSmBqAxg5vPj2TmRDmfkzw+cTzAElWljhcU=
 github.com/StackExchange/wmi v1.2.1 h1:VIkavFPXSjcnS+O8yTq7NI32k0R5Aj+v39y29VYDOSA=
 github.com/StackExchange/wmi v1.2.1/go.mod h1:rcmrprowKIVzvc+NUiLncP2uuArMWLCbu9SBzvHz7e8=