feat: 增加内存存储

.env.example

@@ -23,6 +23,9 @@ ENABLE_GZIP=true
 DATABASE_DSN=user:password@tcp(localhost:3306)/gpt_load?charset=utf8mb4&parseTime=True&loc=Local
 DB_AUTO_MIGRATE=true
 
+# Redis配置
+REDIS_DSN=redis://:password@localhost:6379/2
+
 # 日志配置
 LOG_LEVEL=info
 LOG_FORMAT=text

cmd/gpt-load/main.go

@@ -22,6 +22,7 @@ import (
 	"gpt-load/internal/proxy"
 	"gpt-load/internal/router"
 	"gpt-load/internal/services"
+	"gpt-load/internal/store"
 	"gpt-load/internal/types"
 
 	"github.com/sirupsen/logrus"
@@ -76,7 +77,14 @@ func main() {
 	// ---
 
 	// --- Service Initialization ---
-	taskService := services.NewTaskService()
+	// 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)
 

go.mod

@@ -11,6 +11,7 @@ require (
 	github.com/go-sql-driver/mysql v1.8.1
 	github.com/google/uuid v1.6.0
 	github.com/joho/godotenv v1.5.1
+	github.com/redis/go-redis/v9 v9.5.3
 	github.com/sirupsen/logrus v1.9.3
 	gorm.io/datatypes v1.2.1
 	gorm.io/driver/mysql v1.6.0
@@ -21,7 +22,9 @@ require (
 	filippo.io/edwards25519 v1.1.0 // indirect
 	github.com/bytedance/sonic v1.13.3 // indirect
 	github.com/bytedance/sonic/loader v0.2.4 // indirect
+	github.com/cespare/xxhash/v2 v2.2.0 // indirect
 	github.com/cloudwego/base64x v0.1.5 // indirect
+	github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect
 	github.com/gabriel-vasile/mimetype v1.4.9 // indirect
 	github.com/gin-contrib/sse v1.1.0 // indirect
 	github.com/go-playground/locales v0.14.1 // indirect

go.sum

@@ -1,10 +1,16 @@
 filippo.io/edwards25519 v1.1.0 h1:FNf4tywRC1HmFuKW5xopWpigGjJKiJSV0Cqo0cJWDaA=
 filippo.io/edwards25519 v1.1.0/go.mod h1:BxyFTGdWcka3PhytdK4V28tE5sGfRvvvRV7EaN4VDT4=
+github.com/bsm/ginkgo/v2 v2.12.0 h1:Ny8MWAHyOepLGlLKYmXG4IEkioBysk6GpaRTLC8zwWs=
+github.com/bsm/ginkgo/v2 v2.12.0/go.mod h1:SwYbGRRDovPVboqFv0tPTcG1sN61LM1Z4ARdbAV9g4c=
+github.com/bsm/gomega v1.27.10 h1:yeMWxP2pV2fG3FgAODIY8EiRE3dy0aeFYt4l7wh6yKA=
+github.com/bsm/gomega v1.27.10/go.mod h1:JyEr/xRbxbtgWNi8tIEVPUYZ5Dzef52k01W3YH0H+O0=
 github.com/bytedance/sonic v1.13.3 h1:MS8gmaH16Gtirygw7jV91pDCN33NyMrPbN7qiYhEsF0=
 github.com/bytedance/sonic v1.13.3/go.mod h1:o68xyaF9u2gvVBuGHPlUVCy+ZfmNNO5ETf1+KgkJhz4=
 github.com/bytedance/sonic/loader v0.1.1/go.mod h1:ncP89zfokxS5LZrJxl5z0UJcsk4M4yY2JpfqGeCtNLU=
 github.com/bytedance/sonic/loader v0.2.4 h1:ZWCw4stuXUsn1/+zQDqeE7JKP+QO47tz7QCNan80NzY=
 github.com/bytedance/sonic/loader v0.2.4/go.mod h1:N8A3vUdtUebEY2/VQC0MyhYeKUFosQU6FxH2JmUe6VI=
+github.com/cespare/xxhash/v2 v2.2.0 h1:DC2CZ1Ep5Y4k3ZQ899DldepgrayRUGE6BBZ/cd9Cj44=
+github.com/cespare/xxhash/v2 v2.2.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
 github.com/cloudwego/base64x v0.1.5 h1:XPciSp1xaq2VCSt6lF0phncD4koWyULpl5bUxbfCyP4=
 github.com/cloudwego/base64x v0.1.5/go.mod h1:0zlkT4Wn5C6NdauXdJRhSKRlJvmclQ1hhJgA0rcu/8w=
 github.com/cloudwego/iasm v0.2.0/go.mod h1:8rXZaNYT2n95jn+zTI1sDr+IgcD2GVs0nlbbQPiEFhY=
@@ -12,6 +18,8 @@ github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ3
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78=
+github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc=
 github.com/gabriel-vasile/mimetype v1.4.9 h1:5k+WDwEsD9eTLL8Tz3L0VnmVh9QxGjRmjBvAG7U/oYY=
 github.com/gabriel-vasile/mimetype v1.4.9/go.mod h1:WnSQhFKJuBlRyLiKohA/2DtIlPFAbguNaG7QCHcyGok=
 github.com/gin-contrib/gzip v1.2.3 h1:dAhT722RuEG330ce2agAs75z7yB+NKvX/ZM1r8w0u2U=
@@ -84,6 +92,8 @@ github.com/pelletier/go-toml/v2 v2.2.4 h1:mye9XuhQ6gvn5h28+VilKrrPoQVanw5PMw/TB0
 github.com/pelletier/go-toml/v2 v2.2.4/go.mod h1:2gIqNv+qfxSVS7cM2xJQKtLSTLUE9V8t9Stt+h56mCY=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/redis/go-redis/v9 v9.5.3 h1:fOAp1/uJG+ZtcITgZOfYFmTKPE7n4Vclj1wZFgRciUU=
+github.com/redis/go-redis/v9 v9.5.3/go.mod h1:hdY0cQFCN4fnSYT6TkisLufl/4W5UIXyv0b/CLO2V2M=
 github.com/rogpeppe/go-internal v1.8.0 h1:FCbCCtXNOY3UtUuHUYaghJg4y7Fd14rXifAYUAtL9R8=
 github.com/rogpeppe/go-internal v1.8.0/go.mod h1:WmiCO8CzOY8rg0OYDC4/i/2WRWAB6poM+XZ2dLUbcbE=
 github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ=

internal/config/manager.go

@@ -47,6 +47,7 @@ type Config struct {
 	CORS        types.CORSConfig        `json:"cors"`
 	Performance types.PerformanceConfig `json:"performance"`
 	Log         types.LogConfig         `json:"log"`
+	RedisDSN    string                  `json:"redis_dsn"`
 }
 
 // NewManager creates a new configuration manager
@@ -109,6 +110,7 @@ func (m *Manager) ReloadConfig() error {
 			FilePath:      getEnvOrDefault("LOG_FILE_PATH", "logs/app.log"),
 			EnableRequest: parseBoolean(os.Getenv("LOG_ENABLE_REQUEST"), true),
 		},
+		RedisDSN: os.Getenv("REDIS_DSN"),
 	}
 	m.config = config
 
@@ -122,6 +124,11 @@ func (m *Manager) ReloadConfig() error {
 	return nil
 }
 
+// GetConfig returns the raw config struct
+func (m *Manager) GetConfig() *Config {
+	return m.config
+}
+
 // GetAuthConfig returns authentication configuration
 func (m *Manager) GetAuthConfig() types.AuthConfig {
 	return m.config.Auth
@@ -142,6 +149,11 @@ func (m *Manager) GetLogConfig() types.LogConfig {
 	return m.config.Log
 }
 
+// GetRedisDSN returns the Redis DSN string.
+func (m *Manager) GetRedisDSN() string {
+	return m.config.RedisDSN
+}
+
 // GetEffectiveServerConfig returns server configuration merged with system settings
 func (m *Manager) GetEffectiveServerConfig() types.ServerConfig {
 	config := m.config.Server

internal/store/factory.go

@@ -0,0 +1,37 @@
+package store
+
+import (
+	"context"
+	"fmt"
+	"gpt-load/internal/types"
+
+	"github.com/redis/go-redis/v9"
+	"github.com/sirupsen/logrus"
+)
+
+// NewStore creates a new store based on the application configuration.
+// It prioritizes Redis if a DSN is provided, otherwise it falls back to an in-memory store.
+func NewStore(cfg types.ConfigManager) (Store, error) {
+	redisDSN := cfg.GetRedisDSN()
+	// Prioritize Redis if configured
+	if redisDSN != "" {
+		logrus.Info("Redis DSN found, initializing Redis store...")
+		opts, err := redis.ParseURL(redisDSN)
+		if err != nil {
+			return nil, fmt.Errorf("failed to parse redis DSN: %w", err)
+		}
+
+		client := redis.NewClient(opts)
+		// Ping the server to ensure a connection is established.
+		if err := client.Ping(context.Background()).Err(); err != nil {
+			return nil, fmt.Errorf("failed to connect to redis: %w", err)
+		}
+
+		logrus.Info("Successfully connected to Redis.")
+		return NewRedisStore(client), nil
+	}
+
+	// Fallback to in-memory store
+	logrus.Info("Redis DSN not configured, falling back to in-memory store.")
+	return NewMemoryStore(), nil
+}

internal/store/memory.go

@@ -0,0 +1,126 @@
+package store
+
+import (
+	"sync"
+	"time"
+)
+
+// memoryStoreItem holds the value and expiration timestamp for a key.
+type memoryStoreItem struct {
+	value     []byte
+	expiresAt int64 // Unix-nano timestamp. 0 for no expiry.
+}
+
+// MemoryStore is an in-memory key-value store that is safe for concurrent use.
+type MemoryStore struct {
+	mu     sync.RWMutex
+	data   map[string]memoryStoreItem
+	stopCh chan struct{} // Channel to stop the cleanup goroutine
+}
+
+// NewMemoryStore creates and returns a new MemoryStore instance.
+// It also starts a background goroutine to periodically clean up expired keys.
+func NewMemoryStore() *MemoryStore {
+	s := &MemoryStore{
+		data:   make(map[string]memoryStoreItem),
+		stopCh: make(chan struct{}),
+	}
+	go s.cleanupLoop(1 * time.Minute)
+	return s
+}
+
+// Close stops the background cleanup goroutine.
+func (s *MemoryStore) Close() error {
+	close(s.stopCh)
+	return nil
+}
+
+// cleanupLoop periodically iterates through the store and removes expired keys.
+func (s *MemoryStore) cleanupLoop(interval time.Duration) {
+	ticker := time.NewTicker(interval)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ticker.C:
+			s.mu.Lock()
+			now := time.Now().UnixNano()
+			for key, item := range s.data {
+				if item.expiresAt > 0 && now > item.expiresAt {
+					delete(s.data, key)
+				}
+			}
+			s.mu.Unlock()
+		case <-s.stopCh:
+			return
+		}
+	}
+}
+
+// Set stores a key-value pair.
+func (s *MemoryStore) Set(key string, value []byte, ttl time.Duration) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	var expiresAt int64
+	if ttl > 0 {
+		expiresAt = time.Now().UnixNano() + ttl.Nanoseconds()
+	}
+
+	s.data[key] = memoryStoreItem{
+		value:     value,
+		expiresAt: expiresAt,
+	}
+	return nil
+}
+
+// Get retrieves a value by its key.
+func (s *MemoryStore) Get(key string) ([]byte, error) {
+	s.mu.RLock()
+	item, exists := s.data[key]
+	s.mu.RUnlock()
+
+	if !exists {
+		return nil, ErrNotFound
+	}
+
+	// Check for expiration
+	if item.expiresAt > 0 && time.Now().UnixNano() > item.expiresAt {
+		// Lazy deletion
+		s.mu.Lock()
+		delete(s.data, key)
+		s.mu.Unlock()
+		return nil, ErrNotFound
+	}
+
+	return item.value, nil
+}
+
+// Delete removes a value by its key.
+func (s *MemoryStore) Delete(key string) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	delete(s.data, key)
+	return nil
+}
+
+// Exists checks if a key exists.
+func (s *MemoryStore) Exists(key string) (bool, error) {
+	s.mu.RLock()
+	item, exists := s.data[key]
+	s.mu.RUnlock()
+
+	if !exists {
+		return false, nil
+	}
+
+	if item.expiresAt > 0 && time.Now().UnixNano() > item.expiresAt {
+		// Lazy deletion
+		s.mu.Lock()
+		delete(s.data, key)
+		s.mu.Unlock()
+		return false, nil
+	}
+
+	return true, nil
+}

internal/store/redis.go

@@ -0,0 +1,55 @@
+package store
+
+import (
+	"context"
+	"errors"
+	"time"
+
+	"github.com/redis/go-redis/v9"
+)
+
+// RedisStore is a Redis-backed key-value store.
+type RedisStore struct {
+	client *redis.Client
+}
+
+// NewRedisStore creates a new RedisStore instance.
+func NewRedisStore(client *redis.Client) *RedisStore {
+	return &RedisStore{client: client}
+}
+
+// Set stores a key-value pair in Redis.
+func (s *RedisStore) Set(key string, value []byte, ttl time.Duration) error {
+	return s.client.Set(context.Background(), key, value, ttl).Err()
+}
+
+// Get retrieves a value from Redis.
+func (s *RedisStore) Get(key string) ([]byte, error) {
+	val, err := s.client.Get(context.Background(), key).Bytes()
+	if err != nil {
+		if errors.Is(err, redis.Nil) {
+			return nil, ErrNotFound
+		}
+		return nil, err
+	}
+	return val, nil
+}
+
+// Delete removes a value from Redis.
+func (s *RedisStore) Delete(key string) error {
+	return s.client.Del(context.Background(), key).Err()
+}
+
+// Exists checks if a key exists in Redis.
+func (s *RedisStore) Exists(key string) (bool, error) {
+	val, err := s.client.Exists(context.Background(), key).Result()
+	if err != nil {
+		return false, err
+	}
+	return val > 0, nil
+}
+
+// Close closes the Redis client connection.
+func (s *RedisStore) Close() error {
+	return s.client.Close()
+}

internal/store/store.go

@@ -0,0 +1,33 @@
+package store
+
+import (
+	"errors"
+	"time"
+)
+
+// ErrNotFound is the error returned when a key is not found in the store.
+var ErrNotFound = errors.New("store: key not found")
+
+// Store is a generic key-value store interface.
+// Implementations of this interface must be safe for concurrent use.
+type Store interface {
+	// Set stores a key-value pair with an optional TTL.
+	// - key: The key (string).
+	// - value: The value ([]byte).
+	// - ttl: The expiration time. If ttl is 0, the key never expires.
+	Set(key string, value []byte, ttl time.Duration) error
+
+	// Get retrieves a value by its key.
+	// It must return store.ErrNotFound if the key does not exist.
+	Get(key string) ([]byte, error)
+
+	// Delete removes a value by its key.
+	// If the key does not exist, this operation should be considered successful (idempotent) and not return an error.
+	Delete(key string) error
+
+	// Exists checks if a key exists in the store.
+	Exists(key string) (bool, error)
+
+	// Close closes the store and releases any underlying resources.
+	Close() error
+}

internal/types/types.go

@@ -12,6 +12,7 @@ type ConfigManager interface {
 	GetPerformanceConfig() PerformanceConfig
 	GetLogConfig() LogConfig
 	GetEffectiveServerConfig() ServerConfig
+	GetRedisDSN() string
 	Validate() error
 	DisplayConfig()
 	ReloadConfig() error