chore: code style (#44)

* some code refactor * service * chore: code style --------- Co-authored-by: naiba <[email protected]>
2024-07-23 23:20:55 +08:00
parent a99bf7066b
commit c0288971d8
5 changed files with 57 additions and 67 deletions
@@ -83,17 +83,14 @@ func GetHost() *model.Host {
 	} else {
 		if hi.VirtualizationRole == "guest" {
 			cpuType = "Virtual"
+			ret.Virtualization = hi.VirtualizationSystem
 		} else {
 			cpuType = "Physical"
+			ret.Virtualization = ""
 		}
 		ret.Platform = hi.Platform
 		ret.PlatformVersion = hi.PlatformVersion
 		ret.Arch = hi.KernelArch
-		if cpuType == "Physical" {
-			ret.Virtualization = ""
-		} else {
-			ret.Virtualization = hi.VirtualizationSystem
-		}
 		ret.BootTime = hi.BootTime
 	}

@@ -218,40 +215,6 @@ func GetState(skipConnectionCount bool, skipProcsCount bool) *model.HostState {
 		}
 	}

-	var tcpConnCount, udpConnCount uint64
-	if !skipConnectionCount {
-		ss_err := true
-		if runtime.GOOS == "linux" {
-			tcpStat, err_tcp := goss.ConnectionsWithProtocol(goss.AF_INET, syscall.IPPROTO_TCP)
-			udpStat, err_udp := goss.ConnectionsWithProtocol(goss.AF_INET, syscall.IPPROTO_UDP)
-			if err_tcp == nil && err_udp == nil {
-				ss_err = false
-				tcpConnCount = uint64(len(tcpStat))
-				udpConnCount = uint64(len(udpStat))
-			}
-			if strings.Contains(CachedIP, ":") {
-				tcpStat6, err_tcp := goss.ConnectionsWithProtocol(goss.AF_INET6, syscall.IPPROTO_TCP)
-				udpStat6, err_udp := goss.ConnectionsWithProtocol(goss.AF_INET6, syscall.IPPROTO_UDP)
-				if err_tcp == nil && err_udp == nil {
-					ss_err = false
-					tcpConnCount += uint64(len(tcpStat6))
-					udpConnCount += uint64(len(udpStat6))
-				}
-			}
-		}
-		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++
-				}
-			}
-		}
-	}
-
 	if agentConfig.Temperature {
 		go updateTemperatureStat()
 		ret.Temperatures = temperatureStat
@@ -265,7 +228,7 @@ func GetState(skipConnectionCount bool, skipProcsCount bool) *model.HostState {
 	ret.NetInTransfer, ret.NetOutTransfer = netInTransfer, netOutTransfer
 	ret.NetInSpeed, ret.NetOutSpeed = netInSpeed, netOutSpeed
 	ret.Uptime = uint64(time.Since(cachedBootTime).Seconds())
-	ret.TcpConnCount, ret.UdpConnCount = tcpConnCount, udpConnCount
+	ret.TcpConnCount, ret.UdpConnCount = getConns(skipConnectionCount)

 	return &ret
 }
@@ -352,6 +315,42 @@ func getDiskTotalAndUsed() (total uint64, used uint64) {
 	return
 }

+func getConns(skipConnectionCount bool) (tcpConnCount, udpConnCount uint64) {
+	if !skipConnectionCount {
+		ss_err := true
+		if runtime.GOOS == "linux" {
+			tcpStat, err_tcp := goss.ConnectionsWithProtocol(goss.AF_INET, syscall.IPPROTO_TCP)
+			udpStat, err_udp := goss.ConnectionsWithProtocol(goss.AF_INET, syscall.IPPROTO_UDP)
+			if err_tcp == nil && err_udp == nil {
+				ss_err = false
+				tcpConnCount = uint64(len(tcpStat))
+				udpConnCount = uint64(len(udpStat))
+			}
+			if strings.Contains(CachedIP, ":") {
+				tcpStat6, err_tcp := goss.ConnectionsWithProtocol(goss.AF_INET6, syscall.IPPROTO_TCP)
+				udpStat6, err_udp := goss.ConnectionsWithProtocol(goss.AF_INET6, syscall.IPPROTO_UDP)
+				if err_tcp == nil && err_udp == nil {
+					ss_err = false
+					tcpConnCount += uint64(len(tcpStat6))
+					udpConnCount += uint64(len(udpStat6))
+				}
+			}
+		}
+		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++
+				}
+			}
+		}
+	}
+	return tcpConnCount, udpConnCount
+}
+
 func updateGPUStat(gpuStat *uint64) {
 	if !atomic.CompareAndSwapInt32(&updateGPUStatus, 0, 1) {
 		return
@@ -49,7 +49,6 @@ var (
 		// "https://freegeoip.app/json/", // 需要 Key
 	}
 	CachedIP, CachedCountry string
-	Sync                    = make(chan bool)
 	httpClientV4            = util.NewSingleStackHTTPClient(time.Second*20, time.Second*5, time.Second*10, false)
 	httpClientV6            = util.NewSingleStackHTTPClient(time.Second*20, time.Second*5, time.Second*10, true)
 )
@@ -58,16 +57,10 @@ var (
 func UpdateIP(useIPv6CountryCode bool, period uint32) {
 	for {
 		util.Println(agentConfig.Debug, "正在更新本地缓存IP信息")
-		var primaryIP, secondaryIP geoIP
-		if useIPv6CountryCode {
-			primaryIP = fetchGeoIP(geoIPApiList, true)
-			secondaryIP = fetchGeoIP(geoIPApiList, false)
-		} else {
-			primaryIP = fetchGeoIP(geoIPApiList, false)
-			secondaryIP = fetchGeoIP(geoIPApiList, true)
-		}
+		ipv4 := fetchGeoIP(geoIPApiList, false)
+		ipv6 := fetchGeoIP(geoIPApiList, true)

-		if primaryIP.IP == "" && secondaryIP.IP == "" {
+		if ipv4.IP == "" && ipv6.IP == "" {
 			if period > 60 {
 				time.Sleep(time.Minute)
 			} else {
@@ -75,21 +68,17 @@ func UpdateIP(useIPv6CountryCode bool, period uint32) {
 			}
 			continue
 		}
-		if primaryIP.IP == "" || secondaryIP.IP == "" {
-			CachedIP = fmt.Sprintf("%s%s", primaryIP.IP, secondaryIP.IP)
+		if ipv4.IP == "" || ipv6.IP == "" {
+			CachedIP = fmt.Sprintf("%s%s", ipv4.IP, ipv6.IP)
 		} else {
-			CachedIP = fmt.Sprintf("%s/%s", primaryIP.IP, secondaryIP.IP)
+			CachedIP = fmt.Sprintf("%s/%s", ipv4.IP, ipv6.IP)
 		}

-		if primaryIP.CountryCode != "" {
-			CachedCountry = primaryIP.CountryCode
-		} else if secondaryIP.CountryCode != "" {
-			CachedCountry = secondaryIP.CountryCode
+		if ipv4.CountryCode != "" {
+			CachedCountry = ipv4.CountryCode
 		}
-
-		select {
-		case Sync <- true:
-		default:
+		if ipv6.CountryCode != "" && (useIPv6CountryCode || CachedCountry == "") {
+			CachedCountry = ipv6.CountryCode
 		}

 		time.Sleep(time.Second * time.Duration(period))