gpt-load/cmd/gpt-load/main.go

// Package main provides the entry point for the GPT-Load proxy server
package main

import (
	"context"
	"fmt"
	"io"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"syscall"
	"time"

	"gpt-load/internal/config"
	"gpt-load/internal/handler"
	"gpt-load/internal/keymanager"
	"gpt-load/internal/middleware"
	"gpt-load/internal/proxy"
	"gpt-load/pkg/types"

	"github.com/gin-gonic/gin"
	"github.com/sirupsen/logrus"
)

func main() {
	// Load configuration
	configManager, err := config.NewManager()
	if err != nil {
		logrus.Fatalf("Failed to load configuration: %v", err)
	}

	// Setup logger
	setupLogger(configManager)

	// Display startup information
	displayStartupInfo(configManager)

	// Create key manager
	keyManager, err := keymanager.NewManager(configManager.GetKeysConfig())
	if err != nil {
		logrus.Fatalf("Failed to create key manager: %v", err)
	}
	defer keyManager.Close()

	// Create proxy server
	proxyServer, err := proxy.NewProxyServer(keyManager, configManager)
	if err != nil {
		logrus.Fatalf("Failed to create proxy server: %v", err)
	}
	defer proxyServer.Close()

	// Create handlers
	handlers := handler.NewHandler(keyManager, configManager)

	// Setup routes
	router := setupRoutes(handlers, proxyServer, configManager)

	// Create HTTP server with optimized timeout configuration
	serverConfig := configManager.GetServerConfig()
	server := &http.Server{
		Addr:           fmt.Sprintf("%s:%d", serverConfig.Host, serverConfig.Port),
		Handler:        router,
		ReadTimeout:    60 * time.Second,  // Increased read timeout for large file uploads
		WriteTimeout:   300 * time.Second, // Increased write timeout for streaming responses
		IdleTimeout:    120 * time.Second, // Increased idle timeout for connection reuse
		MaxHeaderBytes: 1 << 20,           // 1MB header limit
	}

	// 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.Infof("Reset keys: http://%s:%d/reset-keys", serverConfig.Host, serverConfig.Port)
		logrus.Infof("Blacklist query: http://%s:%d/blacklist", serverConfig.Host, serverConfig.Port)
		logrus.Info("")

		if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
			logrus.Fatalf("Server startup failed: %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...")

	// Give outstanding requests a deadline for completion
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()

	// Attempt graceful shutdown
	if err := server.Shutdown(ctx); err != nil {
		logrus.Errorf("Server forced to shutdown: %v", err)
	} else {
		logrus.Info("Server exited gracefully")
	}
}

// setupRoutes configures the HTTP routes
func setupRoutes(handlers *handler.Handler, proxyServer *proxy.ProxyServer, configManager types.ConfigManager) *gin.Engine {
	// Set Gin mode
	gin.SetMode(gin.ReleaseMode)

	router := gin.New()

	// Add middleware
	router.Use(middleware.Recovery())
	router.Use(middleware.Logger(configManager.GetLogConfig()))
	router.Use(middleware.CORS(configManager.GetCORSConfig()))
	router.Use(middleware.RateLimiter(configManager.GetPerformanceConfig()))

	// Add authentication middleware if enabled
	if configManager.GetAuthConfig().Enabled {
		router.Use(middleware.Auth(configManager.GetAuthConfig()))
	}

	// Management endpoints
	router.GET("/health", handlers.Health)
	router.GET("/stats", handlers.Stats)
	router.GET("/blacklist", handlers.Blacklist)
	router.GET("/reset-keys", handlers.ResetKeys)
	router.GET("/config", handlers.GetConfig) // Debug endpoint

	// Handle 405 Method Not Allowed
	router.NoMethod(handlers.MethodNotAllowed)

	// Proxy all other requests (this handles 404 as well)
	router.NoRoute(proxyServer.HandleProxy)

	return router
}

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

// displayStartupInfo shows startup information
func displayStartupInfo(configManager types.ConfigManager) {
	serverConfig := configManager.GetServerConfig()
	keysConfig := configManager.GetKeysConfig()
	openaiConfig := configManager.GetOpenAIConfig()
	authConfig := configManager.GetAuthConfig()
	corsConfig := configManager.GetCORSConfig()
	perfConfig := configManager.GetPerformanceConfig()
	logConfig := configManager.GetLogConfig()

	logrus.Info("Current Configuration:")
	logrus.Infof("   Server: %s:%d", serverConfig.Host, serverConfig.Port)
	logrus.Infof("   Keys file: %s", keysConfig.FilePath)
	logrus.Infof("   Start index: %d", keysConfig.StartIndex)
	logrus.Infof("   Blacklist threshold: %d errors", keysConfig.BlacklistThreshold)
	logrus.Infof("   Max retries: %d", keysConfig.MaxRetries)
	logrus.Infof("   Upstream URL: %s", openaiConfig.BaseURL)
	logrus.Infof("   Request timeout: %dms", openaiConfig.Timeout)

	authStatus := "disabled"
	if authConfig.Enabled {
		authStatus = "enabled"
	}
	logrus.Infof("   Authentication: %s", authStatus)

	corsStatus := "disabled"
	if corsConfig.Enabled {
		corsStatus = "enabled"
	}
	logrus.Infof("   CORS: %s", corsStatus)
	logrus.Infof("   Max concurrent requests: %d", perfConfig.MaxConcurrentRequests)

	gzipStatus := "disabled"
	if perfConfig.EnableGzip {
		gzipStatus = "enabled"
	}
	logrus.Infof("   Gzip compression: %s", gzipStatus)

	requestLogStatus := "enabled"
	if !logConfig.EnableRequest {
		requestLogStatus = "disabled"
	}
	logrus.Infof("   Request logging: %s", requestLogStatus)
}