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
	FileCleanupWorker  *service.FileCleanupWorker
	DeviceCleanupWorker *service.DeviceCleanupWorker
	GRPCServer         *runner.Server
	Server             *http.Server
	Publisher          ws.MessagePublisher
	TaskDispatcher     runner.TaskDispatcher
	Logger             *zap.Logger
	PushWorker         *service.PushWorker
}

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

// 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)
	}

	// 2.1.1 聊天文件过期清理
	if a.FileCleanupWorker != nil {
		a.FileCleanupWorker.Start(ctx)
	}

	// 2.1.2 设备 registration_id 清理（清理不活跃设备 token）
	if a.DeviceCleanupWorker != nil {
		a.DeviceCleanupWorker.Start(ctx)
	}

	// 2.2 启动 PushWorker（Redis Stream 异步推送）
	if a.PushWorker != nil {
		a.PushWorker.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()
	}

	if a.FileCleanupWorker != nil {
		a.Logger.Info("Stopping file cleanup worker")
		a.FileCleanupWorker.Stop()
	}

	if a.DeviceCleanupWorker != nil {
		a.Logger.Info("Stopping device cleanup worker")
		a.DeviceCleanupWorker.Stop()
	}

	// 5.1 停止 PushWorker
	if a.PushWorker != nil {
		a.Logger.Info("Stopping push worker (stream)")
		a.PushWorker.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
}