gpt-load/internal/app/app.go

// Package app provides the main application logic and lifecycle management.
package app

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

	"gpt-load/internal/config"
	"gpt-load/internal/keypool"
	"gpt-load/internal/models"
	"gpt-load/internal/proxy"
	"gpt-load/internal/services"
	"gpt-load/internal/store"
	"gpt-load/internal/types"

	"github.com/gin-gonic/gin"
	"github.com/sirupsen/logrus"
	"go.uber.org/dig"
	"gorm.io/gorm"
)

// App holds all services and manages the application lifecycle.
type App struct {
	engine            *gin.Engine
	configManager     types.ConfigManager
	settingsManager   *config.SystemSettingsManager
	groupManager      *services.GroupManager
	logCleanupService *services.LogCleanupService
	requestLogService *services.RequestLogService
	keyCronService    *services.KeyCronService
	keyValidationPool *services.KeyValidationPool
	keyPoolProvider   *keypool.KeyProvider
	leaderService     *services.LeaderService
	proxyServer       *proxy.ProxyServer
	storage           store.Store
	db                *gorm.DB
	httpServer        *http.Server
	wg                sync.WaitGroup
}

// AppParams defines the dependencies for the App.
type AppParams struct {
	dig.In
	Engine            *gin.Engine
	ConfigManager     types.ConfigManager
	SettingsManager   *config.SystemSettingsManager
	GroupManager      *services.GroupManager
	LogCleanupService *services.LogCleanupService
	RequestLogService *services.RequestLogService
	KeyCronService    *services.KeyCronService
	KeyValidationPool *services.KeyValidationPool
	KeyPoolProvider   *keypool.KeyProvider
	LeaderService     *services.LeaderService
	ProxyServer       *proxy.ProxyServer
	Storage           store.Store
	DB                *gorm.DB
}

// NewApp is the constructor for App, with dependencies injected by dig.
func NewApp(params AppParams) *App {
	return &App{
		engine:            params.Engine,
		configManager:     params.ConfigManager,
		settingsManager:   params.SettingsManager,
		groupManager:      params.GroupManager,
		logCleanupService: params.LogCleanupService,
		requestLogService: params.RequestLogService,
		keyCronService:    params.KeyCronService,
		keyValidationPool: params.KeyValidationPool,
		keyPoolProvider:   params.KeyPoolProvider,
		leaderService:     params.LeaderService,
		proxyServer:       params.ProxyServer,
		storage:           params.Storage,
		db:                params.DB,
	}
}

// Start runs the application, it is a non-blocking call.
func (a *App) Start() error {

	// 启动 Leader Service 并等待选举结果
	if err := a.leaderService.Start(); err != nil {
		return fmt.Errorf("leader service failed to start: %w", err)
	}

	// Leader 节点执行初始化，Follower 节点等待
	if a.leaderService.IsLeader() {
		logrus.Info("Leader mode. Performing initial one-time tasks...")
		acquired, err := a.leaderService.AcquireInitializingLock()
		if err != nil {
			return fmt.Errorf("failed to acquire initializing lock: %w", err)
		}
		if !acquired {
			logrus.Warn("Could not acquire initializing lock, another leader might be active. Switching to follower mode for initialization.")
			if err := a.leaderService.WaitForInitializationToComplete(); err != nil {
				return fmt.Errorf("failed to wait for initialization as a fallback follower: %w", err)
			}
		} else {
			defer a.leaderService.ReleaseInitializingLock()

			// 数据库迁移
			if err := a.db.AutoMigrate(
				&models.RequestLog{},
				&models.GroupHourlyStat{},
				&models.APIKey{},
				&models.SystemSetting{},
				&models.Group{},
			); err != nil {
				return fmt.Errorf("database auto-migration failed: %w", err)
			}
			logrus.Info("Database auto-migration completed.")

			// 初始化系统设置
			if err := a.settingsManager.EnsureSettingsInitialized(); err != nil {
				return fmt.Errorf("failed to initialize system settings: %w", err)
			}
			logrus.Info("System settings initialized in DB.")

			a.settingsManager.Initialize(a.storage, a.groupManager, a.leaderService)

			// 从数据库加载密钥到 Redis
			if err := a.keyPoolProvider.LoadKeysFromDB(); err != nil {
				return fmt.Errorf("failed to load keys into key pool: %w", err)
			}
			logrus.Debug("API keys loaded into Redis cache by leader.")
		}
	} else {
		logrus.Info("Follower Mode. Waiting for leader to complete initialization.")
		if err := a.leaderService.WaitForInitializationToComplete(); err != nil {
			return fmt.Errorf("follower failed to start: %w", err)
		}
		a.settingsManager.Initialize(a.storage, a.groupManager, a.leaderService)
	}

	// 显示配置并启动所有后台服务
	a.configManager.DisplayServerConfig()

	a.groupManager.Initialize()

	a.requestLogService.Start()
	a.logCleanupService.Start()
	a.keyValidationPool.Start()
	a.keyCronService.Start()

	// Create HTTP server
	serverConfig := a.configManager.GetEffectiveServerConfig()
	a.httpServer = &http.Server{
		Addr:           fmt.Sprintf("%s:%d", serverConfig.Host, serverConfig.Port),
		Handler:        a.engine,
		ReadTimeout:    time.Duration(serverConfig.ReadTimeout) * time.Second,
		WriteTimeout:   time.Duration(serverConfig.WriteTimeout) * time.Second,
		IdleTimeout:    time.Duration(serverConfig.IdleTimeout) * time.Second,
		MaxHeaderBytes: 1 << 20,
	}

	// Start HTTP server in a new goroutine
	go func() {
		logrus.Info("GPT-Load proxy server started successfully")
		logrus.Infof("Server address: http://%s:%d", serverConfig.Host, serverConfig.Port)
		logrus.Info("")
		if err := a.httpServer.ListenAndServe(); err != nil && err != http.ErrServerClosed {
			logrus.Fatalf("Server startup failed: %v", err)
		}
	}()

	return nil
}

// Stop gracefully shuts down the application.
func (a *App) Stop(ctx context.Context) {
	logrus.Info("Shutting down server...")

	// Shutdown http server
	if err := a.httpServer.Shutdown(ctx); err != nil {
		logrus.Errorf("Server forced to shutdown: %v", err)
	}

	// Stop background services
	a.keyCronService.Stop()
	a.keyValidationPool.Stop()
	a.leaderService.Stop()
	a.logCleanupService.Stop()
	a.requestLogService.Stop()
	a.groupManager.Stop()
	a.settingsManager.Stop()
	a.storage.Close()

	logrus.Info("Server exited gracefully")
}