From 960266bf71ba5db83f4aecffb5a2c036f3e8a739 Mon Sep 17 00:00:00 2001
From: naiba <hi@nai.ba>
Date: Mon, 16 Aug 2021 23:26:38 +0800
Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20[agent]=20=E5=8F=AF=E5=85=B3?=
 =?UTF-8?q?=E9=97=AD=E7=9A=84=E8=BF=9B=E7=A8=8B=E6=95=B0=E8=BF=9E=E6=8E=A5?=
 =?UTF-8?q?=E6=95=B0=E7=9B=91=E6=8E=A7?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 README.md                    |  7 +++++
 cmd/agent/main.go            | 11 +++----
 cmd/agent/monitor/monitor.go | 57 ++++++++++++++++++++++++------------
 3 files changed, 52 insertions(+), 23 deletions(-)

diff --git a/README.md b/README.md
index 0e0470d..6d70f5c 100644
--- a/README.md
+++ b/README.md
@@ -221,6 +221,13 @@ URL 里面也可放置占位符，请求时会进行简单的字符串替换。
 
 </details>
 
+<details>
+    <summary>如何禁用连接数/进程数等占用资源的监控？</summary>
+
+编辑 `/etc/systemd/system/nezha-agent.service`，在 `ExecStart=` 这一行的末尾加上 `-kconn` 就是不监控连接数，加上 `-kprocess` 就是不监控进程数
+
+</details>
+
 <details>
     <summary>Agent 不断重启/无法启动 ？</summary>
 
diff --git a/cmd/agent/main.go b/cmd/agent/main.go
index 9048efa..fd8efed 100644
--- a/cmd/agent/main.go
+++ b/cmd/agent/main.go
@@ -33,10 +33,9 @@ func init() {
 }
 
 var (
-	server       string
-	clientSecret string
-	version      string
-	debug        bool
+	server, clientSecret, version string
+	debug                         bool
+	stateConf                     monitor.GetStateConfig
 )
 
 var (
@@ -67,6 +66,8 @@ func main() {
 	flag.BoolVar(&debug, "d", false, "开启调试信息")
 	flag.StringVar(&server, "s", "localhost:5555", "管理面板RPC端口")
 	flag.StringVar(&clientSecret, "p", "", "Agent连接Secret")
+	flag.BoolVar(&stateConf.SkipConnectionCount, "kconn", false, "不监控连接数")
+	flag.BoolVar(&stateConf.SkipProcessCount, "kprocess", false, "不监控进程数")
 	flag.Parse()
 
 	if server == "" || clientSecret == "" {
@@ -271,7 +272,7 @@ func reportState() {
 		if client != nil && inited {
 			monitor.TrackNetworkSpeed()
 			timeOutCtx, cancel := context.WithTimeout(context.Background(), networkTimeOut)
-			_, err = client.ReportSystemState(timeOutCtx, monitor.GetState().PB())
+			_, err = client.ReportSystemState(timeOutCtx, monitor.GetState(stateConf).PB())
 			cancel()
 			if err != nil {
 				println("reportState error", err)
diff --git a/cmd/agent/monitor/monitor.go b/cmd/agent/monitor/monitor.go
index c64b0ff..13e5a60 100644
--- a/cmd/agent/monitor/monitor.go
+++ b/cmd/agent/monitor/monitor.go
@@ -6,6 +6,7 @@ import (
 	"runtime"
 	"strings"
 	"sync/atomic"
+	"syscall"
 	"time"
 
 	"github.com/shirou/gopsutil/v3/cpu"
@@ -73,7 +74,12 @@ func GetHost() *model.Host {
 	}
 }
 
-func GetState() *model.HostState {
+type GetStateConfig struct {
+	SkipConnectionCount bool
+	SkipProcessCount    bool
+}
+
+func GetState(conf GetStateConfig) *model.HostState {
 	hi, _ := host.Info()
 	mv, _ := mem.VirtualMemory()
 
@@ -94,21 +100,36 @@ func GetState() *model.HostState {
 	_, diskUsed := getDiskTotalAndUsed()
 	loadStat, _ := load.Avg()
 
-	tcpConns, _ := net.Connections("tcp")
-	udpConns, _ := net.Connections("udp")
+	var tcpConnCount, udpConnCount uint64
 
-	ps, _ := process.Pids()
-	// log.Println("pids", len(ps), err)
-	// var threads uint64
-	// for i := 0; i < len(ps); i++ {
-	// 	p, err := process.NewProcess(ps[i])
-	// 	if err != nil {
-	// 		continue
-	// 	}
-	// 	c, _ := p.NumThreads()
-	// 	threads += uint64(c)
-	// }
-	// log.Println("threads", threads)
+	if !conf.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++
+			}
+		}
+	}
+
+	var processCount uint64
+	if !conf.SkipProcessCount {
+		ps, _ := process.Pids()
+		processCount = uint64(len(ps))
+		// log.Println("pids", len(ps), err)
+		// var threads uint64
+		// for i := 0; i < len(ps); i++ {
+		// 	p, err := process.NewProcess(ps[i])
+		// 	if err != nil {
+		// 		continue
+		// 	}
+		// 	c, _ := p.NumThreads()
+		// 	threads += uint64(c)
+		// }
+		// log.Println("threads", threads)
+	}
 
 	return &model.HostState{
 		CPU:            cpuPercent,
@@ -123,9 +144,9 @@ func GetState() *model.HostState {
 		Load1:          loadStat.Load1,
 		Load5:          loadStat.Load5,
 		Load15:         loadStat.Load15,
-		TcpConnCount:   uint64(len(tcpConns)),
-		UdpConnCount:   uint64(len(udpConns)),
-		ProcessCount:   uint64(len(ps)),
+		TcpConnCount:   tcpConnCount,
+		UdpConnCount:   udpConnCount,
+		ProcessCount:   processCount,
 	}
 }