backend/internal/grpc/runner/server.go

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
}