backend/cmd/server/app.go

package main

import (
	"context"
	"fmt"
	"net/http"
	"time"

	"with_you/internal/config"
	"with_you/internal/grpc/runner"
	"with_you/internal/pkg/ws"
	"with_you/internal/router"
	"with_you/internal/service"

	"go.uber.org/zap"
	"gorm.io/gorm"
)

// App 应用程序结构体
type App struct {
	Config          *config.Config
	DB              *gorm.DB
	Router          *router.Router
	PushService     service.PushService
	HotRankWorker   *service.HotRankWorker
	GRPCServer      *runner.Server
	Server          *http.Server
	Publisher       ws.MessagePublisher
	TaskDispatcher  runner.TaskDispatcher
	Logger          *zap.Logger
}

// NewApp 创建应用程序
func NewApp(
	cfg *config.Config,
	db *gorm.DB,
	r *router.Router,
	pushService service.PushService,
	hotRankWorker *service.HotRankWorker,
	grpcServer *runner.Server,
	publisher ws.MessagePublisher,
	taskDispatcher runner.TaskDispatcher,
	logger *zap.Logger,
) *App {
	return &App{
		Config:          cfg,
		DB:              db,
		Router:          r,
		PushService:     pushService,
		HotRankWorker:   hotRankWorker,
		GRPCServer:      grpcServer,
		Publisher:       publisher,
		TaskDispatcher:  taskDispatcher,
		Logger:          logger,
	}
}

// Start 启动应用程序
func (a *App) Start(ctx context.Context) error {
	// 1. 启动 gRPC 服务器（在 HTTP 服务器之前）
	if a.GRPCServer != nil {
		a.Logger.Info("Starting gRPC server")
		if err := a.GRPCServer.Start(); err != nil {
			return fmt.Errorf("failed to start gRPC server: %w", err)
		}
	}

	// 2. 启动推送 Worker
	a.PushService.StartPushWorker(ctx)

	// 2.1 热门榜定时重算
	if a.HotRankWorker != nil {
		a.HotRankWorker.Start(ctx)
	}

	// 3. 创建 HTTP 服务器
	addr := fmt.Sprintf("%s:%d", a.Config.Server.Host, a.Config.Server.Port)
	a.Logger.Info("Starting HTTP server",
		zap.String("address", addr),
	)

	a.Server = &http.Server{
		Addr:              addr,
		Handler:           a.Router.Engine(),
		ReadHeaderTimeout: 10 * time.Second,
		ReadTimeout:       30 * time.Second,
		WriteTimeout:      60 * time.Second,
		IdleTimeout:       120 * time.Second,
	}

	// 4. 启动 HTTP 服务器
	go func() {
		if err := a.Server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
			a.Logger.Fatal("HTTP server listen error",
				zap.Error(err),
			)
		}
	}()

	return nil
}

// Stop 停止应用程序
func (a *App) Stop(ctx context.Context) error {
	a.Logger.Info("Shutting down server")

	// 1. 关闭所有 WebSocket 连接
	if a.Publisher != nil {
		a.Logger.Info("Closing WebSocket connections")
		switch p := a.Publisher.(type) {
		case *ws.Bus:
			p.Close()
		case *ws.Hub:
			p.CloseAllConnections()
		}
	}

	// 3. 停止 HTTP 服务器
	if a.Server != nil {
		a.Logger.Info("Stopping HTTP server")
		if err := a.Server.Shutdown(ctx); err != nil {
			a.Logger.Error("HTTP server forced to shutdown",
				zap.Error(err),
			)
		}
	}

	// 4. 停止 gRPC 服务器
	if a.GRPCServer != nil {
		a.Logger.Info("Stopping gRPC server")
		a.GRPCServer.Stop()
	}

	// 4.1 停止 TaskBus（集群模式）
	if a.TaskDispatcher != nil {
		if bus, ok := a.TaskDispatcher.(*runner.TaskBus); ok {
			a.Logger.Info("Stopping TaskBus")
			bus.Stop()
		}
	}

	// 5. 停止推送 Worker
	a.Logger.Info("Stopping push worker")
	a.PushService.StopPushWorker()

	if a.HotRankWorker != nil {
		a.Logger.Info("Stopping hot rank worker")
		a.HotRankWorker.Stop()
	}

	// 6. 关闭数据库连接
	a.Logger.Info("Closing database connection")
	sqlDB, err := a.DB.DB()
	if err == nil {
		sqlDB.Close()
	}

	a.Logger.Info("Server exited")
	return nil
}