Skip to content

Commit

Permalink
feat: 在面板重启后及时更新host信息
Browse files Browse the repository at this point in the history
  • Loading branch information
naiba committed Aug 2, 2024
1 parent f837a6f commit 1f77062
Show file tree
Hide file tree
Showing 2 changed files with 40 additions and 31 deletions.
68 changes: 38 additions & 30 deletions cmd/agent/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -57,11 +57,11 @@ type AgentCliParam struct {
}

var (
version string
arch string
client pb.NezhaServiceClient
inited bool
resolver = &net.Resolver{PreferGo: true}
version string
arch string
client pb.NezhaServiceClient
initialized bool
resolver = &net.Resolver{PreferGo: true}
)

var agentCmd = &cobra.Command{
Expand Down Expand Up @@ -222,7 +222,7 @@ func run() {
go pty.DownloadDependency()
}
// 上报服务器信息
go reportState()
go reportStateDaemon()
// 更新IP信息
go monitor.UpdateIP(agentCliParam.UseIPv6CountryCode, agentCliParam.IPReportPeriod)

Expand All @@ -240,7 +240,7 @@ func run() {
var conn *grpc.ClientConn

retry := func() {
inited = false
initialized = false
println("Error to close connection ...")
if conn != nil {
conn.Close()
Expand Down Expand Up @@ -280,7 +280,7 @@ func run() {
continue
}
cancel()
inited = true
initialized = true
// 执行 Task
tasks, err := client.RequestTask(context.Background(), monitor.GetHost().PB())
if err != nil {
Expand Down Expand Up @@ -382,7 +382,7 @@ func doTask(task *pb.Task) {
result.Id = task.GetId()
result.Type = task.GetType()
switch task.GetType() {
case model.TaskTypeHTTPGET:
case model.TaskTypeHTTPGet:
handleHttpGetTask(task, &result)
case model.TaskTypeICMPPing:
handleIcmpPingTask(task, &result)
Expand All @@ -398,6 +398,9 @@ func doTask(task *pb.Task) {
case model.TaskTypeNAT:
handleNATTask(task)
return
case model.TaskTypeReportHostInfo:
reportState(time.Time{})
return
case model.TaskTypeKeepalive:
return
default:
Expand All @@ -407,36 +410,41 @@ func doTask(task *pb.Task) {
client.ReportTask(context.Background(), &result)
}

// reportState 向server上报状态信息
func reportState() {
// reportStateDaemon 向server上报状态信息
func reportStateDaemon() {
var lastReportHostInfo time.Time
var err error
defer println("reportState exit", time.Now(), "=>", err)
for {
// 为了更准确的记录时段流量,inited 后再上传状态信息
if client != nil && inited {
monitor.TrackNetworkSpeed()
timeOutCtx, cancel := context.WithTimeout(context.Background(), networkTimeOut)
_, err = client.ReportSystemState(timeOutCtx, monitor.GetState(agentCliParam.SkipConnectionCount, agentCliParam.SkipProcsCount).PB())
cancel()
if err != nil {
println("reportState error", err)
time.Sleep(delayWhenError)
}
// 每10分钟重新获取一次硬件信息
if lastReportHostInfo.Before(time.Now().Add(-10 * time.Minute)) {
lastReportHostInfo = time.Now()
client.ReportSystemInfo(context.Background(), monitor.GetHost().PB())
if monitor.GeoQueryIP != "" {
geoip, err := client.LookupGeoIP(context.Background(), &pb.GeoIP{Ip: monitor.GeoQueryIP})
if err == nil {
monitor.CachedCountryCode = geoip.GetCountryCode()
}
lastReportHostInfo = reportState(lastReportHostInfo)
time.Sleep(time.Second * time.Duration(agentCliParam.ReportDelay))
}
}

func reportState(lastReportHostInfo time.Time) time.Time {
if client != nil && initialized {
monitor.TrackNetworkSpeed()
timeOutCtx, cancel := context.WithTimeout(context.Background(), networkTimeOut)
_, err := client.ReportSystemState(timeOutCtx, monitor.GetState(agentCliParam.SkipConnectionCount, agentCliParam.SkipProcsCount).PB())
cancel()
if err != nil {
println("reportState error", err)
time.Sleep(delayWhenError)
}
// 每10分钟重新获取一次硬件信息
if lastReportHostInfo.Before(time.Now().Add(-10 * time.Minute)) {
lastReportHostInfo = time.Now()
client.ReportSystemInfo(context.Background(), monitor.GetHost().PB())
if monitor.GeoQueryIP != "" {
geoip, err := client.LookupGeoIP(context.Background(), &pb.GeoIP{Ip: monitor.GeoQueryIP})
if err == nil {
monitor.CachedCountryCode = geoip.GetCountryCode()
}
}
}
time.Sleep(time.Second * time.Duration(agentCliParam.ReportDelay))
}
return lastReportHostInfo
}

// doSelfUpdate 执行更新检查 如果更新成功则会结束进程
Expand Down
3 changes: 2 additions & 1 deletion model/task.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ package model

const (
_ = iota
TaskTypeHTTPGET
TaskTypeHTTPGet
TaskTypeICMPPing
TaskTypeTCPPing
TaskTypeCommand
Expand All @@ -11,6 +11,7 @@ const (
TaskTypeKeepalive
TaskTypeTerminalGRPC
TaskTypeNAT
TaskTypeReportHostInfo
)

type TerminalTask struct {
Expand Down

0 comments on commit 1f77062

Please sign in to comment.