Files
gpt-load/internal/app/app.go
2025-07-12 10:15:07 +08:00

192 lines
6.0 KiB
Go
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

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