From b484fe53f0d3c49715daf59254024aecc23aefd7 Mon Sep 17 00:00:00 2001
From: naiba <hi@nai.ba>
Date: Wed, 4 Dec 2024 21:08:28 +0800
Subject: [PATCH] feat: grpc keepalive

---
 cmd/dashboard/main.go    |  1 -
 cmd/dashboard/rpc/rpc.go | 24 +++++++++---------------
 model/service.go         |  2 +-
 3 files changed, 10 insertions(+), 17 deletions(-)

diff --git a/cmd/dashboard/main.go b/cmd/dashboard/main.go
index 673592cb2a..bd8e999b89 100644
--- a/cmd/dashboard/main.go
+++ b/cmd/dashboard/main.go
@@ -120,7 +120,6 @@ func main() {
 	singleton.CleanServiceHistory()
 	serviceSentinelDispatchBus := make(chan model.Service) // 用于传递服务监控任务信息的channel
 	go rpc.DispatchTask(serviceSentinelDispatchBus)
-	go rpc.DispatchKeepalive()
 	go singleton.AlertSentinelStart()
 	singleton.NewServiceSentinel(serviceSentinelDispatchBus)
 
diff --git a/cmd/dashboard/rpc/rpc.go b/cmd/dashboard/rpc/rpc.go
index 2f7ea47a58..1299fbcb49 100644
--- a/cmd/dashboard/rpc/rpc.go
+++ b/cmd/dashboard/rpc/rpc.go
@@ -9,6 +9,7 @@ import (
 	"time"
 
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/keepalive"
 	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/peer"
 
@@ -21,7 +22,14 @@ import (
 )
 
 func ServeRPC() *grpc.Server {
-	server := grpc.NewServer(grpc.ChainUnaryInterceptor(getRealIp, waf))
+	server := grpc.NewServer(grpc.KeepaliveEnforcementPolicy(
+		keepalive.EnforcementPolicy{
+			MinTime:             time.Second * 20,
+			PermitWithoutStream: true,
+		}), grpc.KeepaliveParams(
+		keepalive.ServerParameters{
+			Time: time.Second * 30,
+		}), grpc.ChainUnaryInterceptor(getRealIp, waf))
 	rpcService.NezhaHandlerSingleton = rpcService.NewNezhaHandler()
 	proto.RegisterNezhaServiceServer(server, rpcService.NezhaHandlerSingleton)
 	return server
@@ -116,20 +124,6 @@ func DispatchTask(serviceSentinelDispatchBus <-chan model.Service) {
 	}
 }
 
-func DispatchKeepalive() {
-	singleton.Cron.AddFunc("@every 60s", func() {
-		singleton.SortedServerLock.RLock()
-		defer singleton.SortedServerLock.RUnlock()
-		for i := 0; i < len(singleton.SortedServerList); i++ {
-			if singleton.SortedServerList[i] == nil || singleton.SortedServerList[i].TaskStream == nil {
-				continue
-			}
-
-			singleton.SortedServerList[i].TaskStream.Send(&proto.Task{Type: model.TaskTypeKeepalive})
-		}
-	})
-}
-
 func ServeNAT(w http.ResponseWriter, r *http.Request, natConfig *model.NAT) {
 	singleton.ServerLock.RLock()
 	server := singleton.ServerList[natConfig.ServerID]
diff --git a/model/service.go b/model/service.go
index 61789ef744..94af0e2636 100644
--- a/model/service.go
+++ b/model/service.go
@@ -19,7 +19,7 @@ const (
 	TaskTypeCommand
 	TaskTypeTerminal
 	TaskTypeUpgrade
-	TaskTypeKeepalive
+	TaskTypeKeepaliveDeprecated
 	TaskTypeTerminalGRPC
 	TaskTypeNAT
 	TaskTypeReportHostInfo