// Package app provides the main application logic and lifecycle management. package app import ( "context" "fmt" "net/http" "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 } // 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.SystemSetting{}, &models.Group{}, &models.APIKey{}, &models.RequestLog{}, &models.GroupHourlyStat{}, ); 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") }