package runner import ( "context" "fmt" "net" "time" "with_you/internal/config" "with_you/proto/runner" "go.uber.org/zap" "google.golang.org/grpc" "google.golang.org/grpc/credentials" "google.golang.org/grpc/credentials/insecure" ) // Server gRPC 服务器 type Server struct { cfg *config.GRPCConfig logger *zap.Logger hub *RunnerHub grpcSrv *grpc.Server } // NewServer 创建新的 gRPC 服务器(内部构造 Hub + TaskManager) func NewServer(cfg *config.GRPCConfig, logger *zap.Logger) *Server { taskManager := NewTaskManager(nil, 60*time.Second, logger) hub := NewRunnerHub(taskManager, logger) taskManager.SetDispatcher(hub) return &Server{ cfg: cfg, logger: logger, hub: hub, } } // NewServerWithDeps 使用外部依赖创建 gRPC 服务器(Wire DI 使用) func NewServerWithDeps(cfg *config.GRPCConfig, logger *zap.Logger, hub *RunnerHub) *Server { return &Server{ cfg: cfg, logger: logger, hub: hub, } } // Start 启动 gRPC 服务器 func (s *Server) Start() error { if !s.cfg.Enabled { s.logger.Info("gRPC server is disabled") return nil } lis, err := net.Listen("tcp", s.cfg.Address()) if err != nil { return fmt.Errorf("failed to listen on %s: %w", s.cfg.Address(), err) } // 配置 gRPC 服务器选项 var opts []grpc.ServerOption // 配置 TLS if s.cfg.TLSEnabled { if s.cfg.TLSCertFile == "" || s.cfg.TLSKeyFile == "" { return fmt.Errorf("TLS enabled but cert/key files not specified") } creds, err := credentials.NewServerTLSFromFile(s.cfg.TLSCertFile, s.cfg.TLSKeyFile) if err != nil { return fmt.Errorf("failed to load TLS credentials: %w", err) } opts = append(opts, grpc.Creds(creds)) s.logger.Info("gRPC server TLS enabled") } else { s.logger.Warn("gRPC server running without TLS") } s.grpcSrv = grpc.NewServer(opts...) runner.RegisterRunnerHubServer(s.grpcSrv, s.hub) s.logger.Info("gRPC server starting", zap.String("address", s.cfg.Address())) go func() { if err := s.grpcSrv.Serve(lis); err != nil { s.logger.Error("gRPC server error", zap.Error(err)) } }() // 启动后台任务清理协程 go s.backgroundCleanup() return nil } // backgroundCleanup 后台清理任务 func (s *Server) backgroundCleanup() { ticker := time.NewTicker(30 * time.Second) defer ticker.Stop() for range ticker.C { // 清理超时的 Runner 连接 s.hub.CheckStaleRunners(90 * time.Second) } } // Stop 停止 gRPC 服务器 func (s *Server) Stop() { if s.grpcSrv != nil { s.logger.Info("Stopping gRPC server") s.grpcSrv.GracefulStop() } } // GetHub 获取 RunnerHub func (s *Server) GetHub() *RunnerHub { return s.hub } // GetTaskManager 获取 TaskManager func (s *Server) GetTaskManager() *TaskManager { return s.hub.taskManager } // Shutdown 优雅关闭 func (s *Server) Shutdown(ctx context.Context) error { if s.grpcSrv == nil { return nil } done := make(chan struct{}) go func() { s.grpcSrv.GracefulStop() close(done) }() select { case <-done: s.logger.Info("gRPC server stopped gracefully") return nil case <-ctx.Done(): s.logger.Warn("gRPC server shutdown timeout, forcing stop") s.grpcSrv.Stop() return ctx.Err() } } // GetClientCredentials 获取客户端凭证配置 func GetClientCredentials(cfg *config.GRPCConfig) (credentials.TransportCredentials, error) { if cfg.TLSEnabled { return credentials.NewClientTLSFromFile(cfg.TLSCertFile, "") } return insecure.NewCredentials(), nil }