feat: 添加ping检测

This commit is contained in:
Akizon77
2025-06-21 16:21:30 +08:00
parent 562c8ffa3c
commit d0052ff27d
4 changed files with 100 additions and 37 deletions

View File

@@ -34,6 +34,7 @@ var RootCmd = &cobra.Command{
} }
go server.DoUploadBasicInfoWorks() go server.DoUploadBasicInfoWorks()
for { for {
server.UpdateBasicInfo()
server.EstablishWebSocketConnection() server.EstablishWebSocketConnection()
} }
}, },

View File

@@ -15,10 +15,6 @@ import (
) )
func DoUploadBasicInfoWorks() { func DoUploadBasicInfoWorks() {
err := uploadBasicInfo()
if err != nil {
log.Println("Error uploading basic info:", err)
}
ticker := time.NewTicker(time.Duration(15) * time.Minute) ticker := time.NewTicker(time.Duration(15) * time.Minute)
for range ticker.C { for range ticker.C {
err := uploadBasicInfo() err := uploadBasicInfo()
@@ -27,7 +23,14 @@ func DoUploadBasicInfoWorks() {
} }
} }
} }
func UpdateBasicInfo() {
err := uploadBasicInfo()
if err != nil {
log.Println("Error uploading basic info:", err)
} else {
log.Println("Basic info uploaded successfully")
}
}
func uploadBasicInfo() error { func uploadBasicInfo() error {
cpu := monitoring.Cpu() cpu := monitoring.Cpu()

View File

@@ -2,6 +2,7 @@ package server
import ( import (
"bytes" "bytes"
"context"
"encoding/json" "encoding/json"
"errors" "errors"
"log" "log"
@@ -83,36 +84,88 @@ func uploadTaskResult(taskID, result string, exitCode int, finishedAt time.Time)
} }
} }
} }
func icmpPing(target string, timeout time.Duration) error {
pinger, err := getPinger(target) // resolveIP 解析域名到 IP 地址,排除 DNS 查询时间
func resolveIP(target string) (string, error) {
// 如果已经是 IP 地址,直接返回
if net.ParseIP(target) != nil {
return target, nil
}
// 解析域名到 IP
addrs, err := net.LookupHost(target)
if err != nil || len(addrs) == 0 {
return "", errors.New("failed to resolve target")
}
return addrs[0], nil // 返回第一个解析的 IP
}
func icmpPing(target string, timeout time.Duration) (int64, error) {
// 先解析 IP 地址
ip, err := resolveIP(target)
if err != nil { if err != nil {
return err return 0, err
}
pinger, err := ping.NewPinger(ip)
if err != nil {
return 0, err
} }
pinger.Count = 1 pinger.Count = 1
pinger.Timeout = timeout pinger.Timeout = timeout
pinger.SetPrivileged(true) pinger.SetPrivileged(true)
return pinger.Run() err = pinger.Run()
}
func getPinger(target string) (*ping.Pinger, error) {
return ping.NewPinger(target)
}
func tcpPing(target string, timeout time.Duration) error {
if !strings.Contains(target, ":") {
target += ":80"
}
conn, err := net.DialTimeout("tcp", target, timeout)
if err != nil { if err != nil {
return err return 0, err
}
stats := pinger.Statistics()
if stats.PacketsRecv == 0 {
return 0, errors.New("no packets received")
}
return stats.AvgRtt.Milliseconds(), nil
}
func tcpPing(target string, timeout time.Duration) (int64, error) {
addr := strings.Split(target, ":")
ip, err := resolveIP(addr[0])
if err != nil {
return 0, err
}
port := "80"
if len(addr) > 1 {
port = addr[1]
}
targetAddr := ip + ":" + port
start := time.Now()
conn, err := net.DialTimeout("tcp", targetAddr, timeout)
if err != nil {
return 0, err
} }
defer conn.Close() defer conn.Close()
return nil return time.Since(start).Milliseconds(), nil
} }
func httpPing(target string, timeout time.Duration) (int64, error) { func httpPing(target string, timeout time.Duration) (int64, error) {
client := http.Client{ if !strings.HasPrefix(target, "http://") && !strings.HasPrefix(target, "https://") {
target = "http://" + target
}
client := &http.Client{
Timeout: timeout, Timeout: timeout,
Transport: &http.Transport{
DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
// 在 Dial 之前解析 IP排除 DNS 时间
host, port, err := net.SplitHostPort(addr)
if err != nil {
return nil, err
}
ip, err := resolveIP(host)
if err != nil {
return nil, err
}
return net.DialTimeout(network, ip+":"+port, timeout)
},
},
} }
start := time.Now() start := time.Now()
resp, err := client.Get(target) resp, err := client.Get(target)
@@ -129,33 +182,39 @@ func httpPing(target string, timeout time.Duration) (int64, error) {
func NewPingTask(conn *ws.SafeConn, taskID uint, pingType, pingTarget string) { func NewPingTask(conn *ws.SafeConn, taskID uint, pingType, pingTarget string) {
if taskID == 0 { if taskID == 0 {
log.Printf("Invalid task ID: %d", taskID)
return return
} }
pingResult := 0 pingResult := 0
timeout := 3 * time.Second timeout := 3 * time.Second // 默认超时时间
switch pingType { switch pingType {
case "icmp": case "icmp":
start := time.Now() if latency, err := icmpPing(pingTarget, timeout); err == nil {
if err := icmpPing(pingTarget, timeout); err == nil { pingResult = int(latency)
pingResult = int(time.Since(start).Milliseconds())
} }
case "tcp": case "tcp":
start := time.Now() if latency, err := tcpPing(pingTarget, timeout); err == nil {
if err := tcpPing(pingTarget, timeout); err == nil { pingResult = int(latency)
pingResult = int(time.Since(start).Milliseconds())
} }
case "http": case "http":
if latency, err := httpPing(pingTarget, timeout); err == nil { if latency, err := httpPing(pingTarget, timeout); err == nil {
pingResult = int(latency) pingResult = int(latency)
} }
default: default:
log.Printf("Unsupported ping type: %s", pingType)
return return
} }
payload := map[string]interface{}{ payload := map[string]interface{}{
"type": "ping_result", "type": "ping_result",
"task_id": taskID, "task_id": taskID,
"ping_type": pingType,
"value": pingResult, "value": pingResult,
"finished_at": time.Now(), "finished_at": time.Now(),
} }
_ = conn.WriteJSON(payload) if err := conn.WriteJSON(payload); err != nil {
log.Printf("Failed to write JSON to WebSocket: %v", err)
}
} }

View File

@@ -37,18 +37,18 @@ func (sc *SafeConn) Close() error {
return sc.conn.Close() return sc.conn.Close()
} }
func (sc *SafeConn) ReadMessage() (int, []byte, error) { func (sc *SafeConn) ReadMessage() (int, []byte, error) {
sc.mu.Lock() // sc.mu.Lock()
defer sc.mu.Unlock() // defer sc.mu.Unlock()
return sc.conn.ReadMessage() return sc.conn.ReadMessage()
} }
func (sc *SafeConn) ReadJSON(v interface{}) error { func (sc *SafeConn) ReadJSON(v interface{}) error {
sc.mu.Lock() // sc.mu.Lock()
defer sc.mu.Unlock() // defer sc.mu.Unlock()
return sc.conn.ReadJSON(v) return sc.conn.ReadJSON(v)
} }
func (sc *SafeConn) SetReadDeadline(t time.Time) error { func (sc *SafeConn) SetReadDeadline(t time.Time) error {
sc.mu.Lock() // sc.mu.Lock()
defer sc.mu.Unlock() // defer sc.mu.Unlock()
return sc.conn.SetReadDeadline(t) return sc.conn.SetReadDeadline(t)
} }
func (sc *SafeConn) GetConn() *websocket.Conn { func (sc *SafeConn) GetConn() *websocket.Conn {