feat: 引入dig服务依赖管理
This commit is contained in:
tbphp
@@ -4,29 +4,18 @@ package main
|
||||
import (
|
||||
"context"
|
||||
"embed"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"os/signal"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
"gpt-load/internal/channel"
|
||||
"gpt-load/internal/app"
|
||||
"gpt-load/internal/config"
|
||||
"gpt-load/internal/db"
|
||||
"gpt-load/internal/handler"
|
||||
"gpt-load/internal/container"
|
||||
"gpt-load/internal/models"
|
||||
"gpt-load/internal/proxy"
|
||||
"gpt-load/internal/router"
|
||||
"gpt-load/internal/services"
|
||||
"gpt-load/internal/store"
|
||||
"gpt-load/internal/types"
|
||||
|
||||
"github.com/sirupsen/logrus"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
//go:embed dist
|
||||
@@ -36,217 +25,53 @@ var buildFS embed.FS
|
||||
var indexPage []byte
|
||||
|
||||
func main() {
|
||||
// Load configuration
|
||||
configManager, err := config.NewManager()
|
||||
// Build the dependency injection container
|
||||
container, err := container.BuildContainer()
|
||||
if err != nil {
|
||||
logrus.Fatalf("Failed to load configuration: %v", err)
|
||||
logrus.Fatalf("Failed to build container: %v", err)
|
||||
}
|
||||
|
||||
// Setup logger
|
||||
setupLogger(configManager)
|
||||
|
||||
// Initialize database
|
||||
database, err := db.InitDB(configManager)
|
||||
if err != nil {
|
||||
logrus.Fatalf("Failed to initialize database: %v", err)
|
||||
// Provide UI assets to the container
|
||||
if err := container.Provide(func() embed.FS { return buildFS }); err != nil {
|
||||
logrus.Fatalf("Failed to provide buildFS: %v", err)
|
||||
}
|
||||
if err := container.Provide(func() []byte { return indexPage }); err != nil {
|
||||
logrus.Fatalf("Failed to provide indexPage: %v", err)
|
||||
}
|
||||
|
||||
// Initialize system settings after database is ready
|
||||
settingsManager := config.GetSystemSettingsManager()
|
||||
if err := settingsManager.InitializeSystemSettings(); err != nil {
|
||||
logrus.Fatalf("Failed to initialize system settings: %v", err)
|
||||
}
|
||||
logrus.Info("System settings initialized")
|
||||
|
||||
// Display current system settings
|
||||
settingsManager.DisplayCurrentSettings()
|
||||
|
||||
// Display startup information
|
||||
configManager.DisplayConfig()
|
||||
|
||||
// Start log cleanup service
|
||||
logCleanupService := services.NewLogCleanupService()
|
||||
logCleanupService.Start()
|
||||
defer logCleanupService.Stop()
|
||||
|
||||
// --- Asynchronous Request Logging Setup ---
|
||||
// Provide the request log channel as a value
|
||||
requestLogChan := make(chan models.RequestLog, 1000)
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(1)
|
||||
go startRequestLogger(database, requestLogChan, &wg)
|
||||
// ---
|
||||
|
||||
// --- Service Initialization ---
|
||||
// Initialize the store first, as other services depend on it.
|
||||
storage, err := store.NewStore(configManager)
|
||||
if err != nil {
|
||||
logrus.Fatalf("Failed to initialize store: %v", err)
|
||||
}
|
||||
defer storage.Close()
|
||||
|
||||
taskService := services.NewTaskService(storage)
|
||||
channelFactory := channel.NewFactory(settingsManager)
|
||||
keyValidatorService := services.NewKeyValidatorService(database, channelFactory, settingsManager)
|
||||
|
||||
// --- Global Key Validation Pool ---
|
||||
KeyValidationPool := services.NewKeyValidationPool(keyValidatorService, configManager)
|
||||
KeyValidationPool.Start()
|
||||
defer KeyValidationPool.Stop()
|
||||
// ---
|
||||
|
||||
keyManualValidationService := services.NewKeyManualValidationService(database, keyValidatorService, taskService, settingsManager, configManager)
|
||||
keyCronService := services.NewKeyCronService(database, settingsManager, KeyValidationPool, storage)
|
||||
keyCronService.Start()
|
||||
defer keyCronService.Stop()
|
||||
|
||||
keyService := services.NewKeyService(database)
|
||||
// ---
|
||||
|
||||
// Create proxy server
|
||||
proxyServer, err := proxy.NewProxyServer(database, channelFactory, requestLogChan)
|
||||
if err != nil {
|
||||
logrus.Fatalf("Failed to create proxy server: %v", err)
|
||||
}
|
||||
defer proxyServer.Close()
|
||||
|
||||
// Create handlers
|
||||
serverHandler := handler.NewServer(database, configManager, settingsManager, keyValidatorService, keyManualValidationService, taskService, keyService)
|
||||
logCleanupHandler := handler.NewLogCleanupHandler(logCleanupService)
|
||||
|
||||
// Setup routes using the new router package
|
||||
appRouter := router.New(serverHandler, proxyServer, logCleanupHandler, configManager, buildFS, indexPage)
|
||||
|
||||
// Create HTTP server with optimized timeout configuration
|
||||
serverConfig := configManager.GetEffectiveServerConfig()
|
||||
server := &http.Server{
|
||||
Addr: fmt.Sprintf("%s:%d", serverConfig.Host, serverConfig.Port),
|
||||
Handler: appRouter,
|
||||
ReadTimeout: time.Duration(serverConfig.ReadTimeout) * time.Second,
|
||||
WriteTimeout: time.Duration(serverConfig.WriteTimeout) * time.Second,
|
||||
IdleTimeout: time.Duration(serverConfig.IdleTimeout) * time.Second,
|
||||
MaxHeaderBytes: 1 << 20, // 1MB header limit
|
||||
if err := container.Provide(func() chan models.RequestLog { return requestLogChan }); err != nil {
|
||||
logrus.Fatalf("Failed to provide request log channel: %v", err)
|
||||
}
|
||||
|
||||
// Start server
|
||||
go func() {
|
||||
logrus.Info("GPT-Load proxy server started successfully")
|
||||
logrus.Infof("Server address: http://%s:%d", serverConfig.Host, serverConfig.Port)
|
||||
logrus.Infof("Statistics: http://%s:%d/stats", serverConfig.Host, serverConfig.Port)
|
||||
logrus.Infof("Health check: http://%s:%d/health", serverConfig.Host, serverConfig.Port)
|
||||
logrus.Info("")
|
||||
// Initialzie global logger
|
||||
if err := container.Invoke(func(configManager types.ConfigManager) {
|
||||
config.SetupLogger(configManager)
|
||||
}); err != nil {
|
||||
logrus.Fatalf("Failed to setup logger: %v", err)
|
||||
}
|
||||
|
||||
if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
|
||||
logrus.Fatalf("Server startup failed: %v", err)
|
||||
// Create and run the application
|
||||
if err := container.Invoke(func(application *app.App, configManager types.ConfigManager) {
|
||||
if err := application.Start(); err != nil {
|
||||
logrus.Fatalf("Failed to start application: %v", err)
|
||||
}
|
||||
}()
|
||||
|
||||
// Wait for interrupt signal to gracefully shutdown the server
|
||||
quit := make(chan os.Signal, 1)
|
||||
signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
|
||||
<-quit
|
||||
logrus.Info("Shutting down server...")
|
||||
// Wait for interrupt signal for graceful shutdown
|
||||
quit := make(chan os.Signal, 1)
|
||||
signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
|
||||
<-quit
|
||||
|
||||
// Give outstanding requests a deadline for completion
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Duration(serverConfig.GracefulShutdownTimeout)*time.Second)
|
||||
defer cancel()
|
||||
// Create a context with timeout for shutdown
|
||||
serverConfig := configManager.GetEffectiveServerConfig()
|
||||
shutdownCtx, cancel := context.WithTimeout(context.Background(), time.Duration(serverConfig.GracefulShutdownTimeout)*time.Second)
|
||||
defer cancel()
|
||||
|
||||
// Attempt graceful shutdown
|
||||
if err := server.Shutdown(ctx); err != nil {
|
||||
logrus.Errorf("Server forced to shutdown: %v", err)
|
||||
}
|
||||
// Perform graceful shutdown
|
||||
application.Stop(shutdownCtx)
|
||||
|
||||
// Close the request log channel and wait for the logger to finish
|
||||
logrus.Info("Closing request log channel...")
|
||||
close(requestLogChan)
|
||||
wg.Wait()
|
||||
logrus.Info("All logs have been written.")
|
||||
|
||||
logrus.Info("Server exited gracefully")
|
||||
}
|
||||
|
||||
// setupLogger, displayStartupInfo, and startRequestLogger functions remain unchanged.
|
||||
// The old setupRoutes and ServeUI functions are now removed from this file.
|
||||
|
||||
// setupLogger configures the logging system
|
||||
func setupLogger(configManager types.ConfigManager) {
|
||||
logConfig := configManager.GetLogConfig()
|
||||
|
||||
// Set log level
|
||||
level, err := logrus.ParseLevel(logConfig.Level)
|
||||
if err != nil {
|
||||
logrus.Warn("Invalid log level, using info")
|
||||
level = logrus.InfoLevel
|
||||
}
|
||||
logrus.SetLevel(level)
|
||||
|
||||
// Set log format
|
||||
if logConfig.Format == "json" {
|
||||
logrus.SetFormatter(&logrus.JSONFormatter{
|
||||
TimestampFormat: time.RFC3339,
|
||||
})
|
||||
} else {
|
||||
logrus.SetFormatter(&logrus.TextFormatter{
|
||||
ForceColors: true,
|
||||
FullTimestamp: true,
|
||||
TimestampFormat: "2006-01-02 15:04:05",
|
||||
})
|
||||
}
|
||||
|
||||
// Setup file logging if enabled
|
||||
if logConfig.EnableFile {
|
||||
// Create log directory if it doesn't exist
|
||||
logDir := filepath.Dir(logConfig.FilePath)
|
||||
if err := os.MkdirAll(logDir, 0755); err != nil {
|
||||
logrus.Warnf("Failed to create log directory: %v", err)
|
||||
} else {
|
||||
// Open log file
|
||||
logFile, err := os.OpenFile(logConfig.FilePath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0666)
|
||||
if err != nil {
|
||||
logrus.Warnf("Failed to open log file: %v", err)
|
||||
} else {
|
||||
// Use both file and stdout
|
||||
logrus.SetOutput(io.MultiWriter(os.Stdout, logFile))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// startRequestLogger runs a background goroutine to batch-insert request logs.
|
||||
func startRequestLogger(db *gorm.DB, logChan <-chan models.RequestLog, wg *sync.WaitGroup) {
|
||||
defer wg.Done()
|
||||
ticker := time.NewTicker(5 * time.Second)
|
||||
defer ticker.Stop()
|
||||
|
||||
logBuffer := make([]models.RequestLog, 0, 100)
|
||||
|
||||
for {
|
||||
select {
|
||||
case logEntry, ok := <-logChan:
|
||||
if !ok {
|
||||
// Channel closed, flush remaining logs and exit
|
||||
if len(logBuffer) > 0 {
|
||||
if err := db.Create(&logBuffer).Error; err != nil {
|
||||
logrus.Errorf("Failed to write remaining request logs: %v", err)
|
||||
}
|
||||
}
|
||||
logrus.Info("Request logger stopped.")
|
||||
return
|
||||
}
|
||||
logBuffer = append(logBuffer, logEntry)
|
||||
if len(logBuffer) >= 100 {
|
||||
if err := db.Create(&logBuffer).Error; err != nil {
|
||||
logrus.Errorf("Failed to write request logs: %v", err)
|
||||
}
|
||||
logBuffer = make([]models.RequestLog, 0, 100) // Reset buffer
|
||||
}
|
||||
case <-ticker.C:
|
||||
// Flush logs periodically
|
||||
if len(logBuffer) > 0 {
|
||||
if err := db.Create(&logBuffer).Error; err != nil {
|
||||
logrus.Errorf("Failed to write request logs on tick: %v", err)
|
||||
}
|
||||
logBuffer = make([]models.RequestLog, 0, 100) // Reset buffer
|
||||
}
|
||||
}
|
||||
}); err != nil {
|
||||
logrus.Fatalf("Failed to run application: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user