Update.

Author: xOS

service/rpc/server.go

@@ -115,10 +115,14 @@ func (s *ServerHandler) RequestTask(h *pb.Host, stream pb.ServerService_RequestT
 	singleton.ServerList[clientID].TaskCloseLock.Unlock()
 	singleton.ServerLock.RUnlock()
 
+	// 创建一个带超时的上下文
+	ctx, cancel := context.WithTimeout(stream.Context(), 30*time.Minute)
+	defer cancel()
+
 	// 监听连接状态，当连接断开时自动清理
 	go func() {
 		select {
-		case <-stream.Context().Done():
+		case <-ctx.Done():
 			// 连接断开时清理资源
 			singleton.ServerLock.RLock()
 			if singleton.ServerList[clientID] != nil {
@@ -134,7 +138,7 @@ func (s *ServerHandler) RequestTask(h *pb.Host, stream pb.ServerService_RequestT
 
 			// 安全地发送关闭信号，使用非阻塞发送
 			select {
-			case closeCh <- stream.Context().Err():
+			case closeCh <- ctx.Err():
 			case <-done: // 如果主goroutine已经退出，停止发送
 				return
 			default:
@@ -146,11 +150,9 @@ func (s *ServerHandler) RequestTask(h *pb.Host, stream pb.ServerService_RequestT
 		}
 	}()
 
-	// 设置超时防止无限等待
-	ctx, cancel := context.WithTimeout(stream.Context(), 30*time.Minute)
-	defer cancel()
 	defer close(done) // 确保监控goroutine被通知停止
 
+	// 等待连接关闭或超时
 	select {
 	case err := <-closeCh:
 		return err

service/singleton/server.go

@@ -1,6 +1,7 @@
 package singleton
 
 import (
+	"fmt"
 	"log"
 	"sort"
 	"sync"
@@ -206,10 +207,30 @@ func CleanupServerState() {
 			server.LastStateBeforeOffline = nil
 			server.TaskStream = nil
 			if server.TaskClose != nil {
+				// 安全关闭通道
+				select {
+				case server.TaskClose <- fmt.Errorf("server state cleanup"):
+				default:
+				}
 				close(server.TaskClose)
 				server.TaskClose = nil
 			}
 		}
+
+		// 清理长时间未活动的连接
+		if server.IsOnline && now.Sub(server.LastActive) > 5*time.Minute {
+			server.TaskCloseLock.Lock()
+			if server.TaskClose != nil {
+				select {
+				case server.TaskClose <- fmt.Errorf("connection timeout"):
+				default:
+				}
+				close(server.TaskClose)
+				server.TaskClose = nil
+			}
+			server.TaskStream = nil
+			server.TaskCloseLock.Unlock()
+		}
 	}
 }
 

service/singleton/singleton.go

@@ -3,6 +3,7 @@ package singleton
 import (
 	"fmt"
 	"log"
+	"runtime"
 	"time"
 
 	"github.com/patrickmn/go-cache"
@@ -66,6 +67,19 @@ func LoadSingleton() {
 	Cron.AddFunc("0 */10 * * * *", func() {
 		CleanMonitorHistory()
 		Cache.DeleteExpired()
+		CleanupServerState() // 添加服务器状态清理
+		SaveAllTrafficToDB() // 保存流量数据到数据库
+	})
+
+	// 添加内存使用监控任务，每5分钟执行一次
+	Cron.AddFunc("0 */5 * * * *", func() {
+		var m runtime.MemStats
+		runtime.ReadMemStats(&m)
+		if m.Alloc > 500*1024*1024 { // 如果内存使用超过500MB
+			log.Printf("内存使用警告: %v MB", m.Alloc/1024/1024)
+			// 触发GC
+			runtime.GC()
+		}
 	})
 }