easyai-ai-gateway/apps/api/cmd/gateway/main.go

package main

import (
	"context"
	"errors"
	"log/slog"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/capacitycontroller"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/httpapi"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)

// @title EasyAI AI Gateway API
// @version 0.1.0
// @description EasyAI AI Gateway 的本地鉴权、平台模型管理、定价、运行策略、钱包和 AI 任务接口。
// @description 受保护接口使用 Authorization: Bearer <JWT 或 API Key>，管理接口只接受 JWT 用户凭证。
// @BasePath /
// @schemes http https
// @securityDefinitions.apikey BearerAuth
// @in header
// @name Authorization
// @description Bearer JWT 或 API Key。
func main() {
	cfg := config.Load()
	logger := slog.New(slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{
		Level: cfg.LogLevel,
	}))
	if err := cfg.Validate(); err != nil {
		logger.Error("invalid gateway configuration", "error", err)
		os.Exit(1)
	}

	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)

	executionMaxConns := cfg.DatabaseMaxConns
	if cfg.DatabaseCriticalMaxConns > 0 {
		executionMaxConns -= cfg.DatabaseCriticalMaxConns
	}
	if cfg.DatabaseRiverMaxConns > 0 {
		executionMaxConns -= cfg.DatabaseRiverMaxConns
	}
	db, err := store.ConnectWithPoolOptions(ctx, cfg.DatabaseURL, store.PostgresPoolOptions{
		MaxConns:                 executionMaxConns,
		MinIdleConns:             cfg.DatabaseMinIdleConns,
		MaxConnIdleTime:          time.Duration(cfg.DatabaseMaxConnIdleSeconds) * time.Second,
		IdleInTransactionTimeout: time.Duration(cfg.DatabaseIdleInTransactionTimeoutSeconds) * time.Second,
		LockTimeout:              time.Duration(cfg.DatabaseLockTimeoutSeconds) * time.Second,
	})
	if err != nil {
		stop()
		logger.Error("connect postgres failed", "error", err)
		os.Exit(1)
	}
	defer db.Close()
	coordinationDB := db
	if cfg.DatabaseCriticalMaxConns > 0 {
		coordinationDB, err = store.ConnectWithPoolOptions(ctx, cfg.DatabaseURL, store.PostgresPoolOptions{
			MaxConns:                 cfg.DatabaseCriticalMaxConns,
			MinIdleConns:             min(cfg.DatabaseCriticalMaxConns, max(cfg.DatabaseMinIdleConns, 1)),
			MaxConnIdleTime:          time.Duration(cfg.DatabaseMaxConnIdleSeconds) * time.Second,
			IdleInTransactionTimeout: time.Duration(cfg.DatabaseIdleInTransactionTimeoutSeconds) * time.Second,
			LockTimeout:              time.Duration(cfg.DatabaseLockTimeoutSeconds) * time.Second,
		})
		if err != nil {
			stop()
			logger.Error("connect critical postgres pool failed", "error", err)
			os.Exit(1)
		}
		defer coordinationDB.Close()
	}
	riverDB := db
	if cfg.DatabaseRiverMaxConns > 0 {
		riverDB, err = store.ConnectWithPoolOptions(ctx, cfg.DatabaseURL, store.PostgresPoolOptions{
			MaxConns:                 cfg.DatabaseRiverMaxConns,
			MinIdleConns:             min(cfg.DatabaseRiverMaxConns, max(cfg.DatabaseMinIdleConns, 1)),
			MaxConnIdleTime:          time.Duration(cfg.DatabaseMaxConnIdleSeconds) * time.Second,
			IdleInTransactionTimeout: time.Duration(cfg.DatabaseIdleInTransactionTimeoutSeconds) * time.Second,
			LockTimeout:              time.Duration(cfg.DatabaseLockTimeoutSeconds) * time.Second,
		})
		if err != nil {
			stop()
			logger.Error("connect River postgres pool failed", "error", err)
			os.Exit(1)
		}
		defer riverDB.Close()
	}
	// Cancel background LISTEN and worker goroutines before closing the pool.
	// This also prevents a startup panic from waiting forever in pgxpool.Close.
	defer stop()
	if cfg.RunsBackgroundWorkers() {
		if recovery, recoveryErr := coordinationDB.RecoverInterruptedRuntimeState(ctx); recoveryErr != nil {
			logger.Error("recover interrupted runtime state failed", "error", recoveryErr)
			os.Exit(1)
		} else if recovery.ReleasedConcurrencyLeases > 0 || recovery.ReleasedRateReservations > 0 || recovery.FailedAttempts > 0 || recovery.FailedTasks > 0 || recovery.RequeuedAsyncTasks > 0 || recovery.CleanedTaskAdmissions > 0 {
			logger.Warn("interrupted runtime state recovered",
				"releasedConcurrencyLeases", recovery.ReleasedConcurrencyLeases,
				"releasedRateReservations", recovery.ReleasedRateReservations,
				"failedAttempts", recovery.FailedAttempts,
				"failedTasks", recovery.FailedTasks,
				"requeuedAsyncTasks", recovery.RequeuedAsyncTasks,
				"cleanedTaskAdmissions", recovery.CleanedTaskAdmissions,
			)
		}
	}

	var handler http.Handler
	if cfg.RunsCapacityController() {
		kubernetes, kubernetesErr := capacitycontroller.NewKubernetesClient(capacitycontroller.KubernetesConfig{
			Namespace: cfg.CapacityControllerNamespace,
			APIServer: cfg.CapacityControllerAPIServer,
			TokenFile: cfg.CapacityControllerTokenFile,
			CAFile:    cfg.CapacityControllerCAFile,
		})
		if kubernetesErr != nil {
			logger.Error("initialize capacity controller Kubernetes client failed", "error", kubernetesErr)
			os.Exit(1)
		}
		controller := capacitycontroller.New(cfg, coordinationDB, kubernetes, logger)
		go controller.Run(ctx)
		handler = controller.Handler()
	} else {
		handler = httpapi.NewServerWithStores(ctx, cfg, db, coordinationDB, riverDB, logger)
	}
	server := &http.Server{
		Addr:              cfg.HTTPAddr,
		Handler:           handler,
		ReadHeaderTimeout: 10 * time.Second,
	}

	go func() {
		logger.Info("easyai ai gateway api started", "addr", cfg.HTTPAddr)
		if err := server.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
			logger.Error("http server failed", "error", err)
			os.Exit(1)
		}
	}()

	<-ctx.Done()
	shutdownCtx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
	defer cancel()
	if err := server.Shutdown(shutdownCtx); err != nil {
		logger.Error("http shutdown failed", "error", err)
		os.Exit(1)
	}
	logger.Info("easyai ai gateway api stopped")
}