From 9b252bbd3ef2dd11832ad54780c809f9bf0a4896 Mon Sep 17 00:00:00 2001 From: tbphp Date: Sat, 5 Jul 2025 15:42:20 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E5=A2=9E=E5=8A=A0=E5=86=85=E5=AD=98?= =?UTF-8?q?=E5=AD=98=E5=82=A8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .env.example | 3 + cmd/gpt-load/main.go | 10 ++- go.mod | 3 + go.sum | 10 +++ internal/config/manager.go | 12 ++++ internal/store/factory.go | 37 +++++++++++ internal/store/memory.go | 126 +++++++++++++++++++++++++++++++++++++ internal/store/redis.go | 55 ++++++++++++++++ internal/store/store.go | 33 ++++++++++ internal/types/types.go | 1 + 10 files changed, 289 insertions(+), 1 deletion(-) create mode 100644 internal/store/factory.go create mode 100644 internal/store/memory.go create mode 100644 internal/store/redis.go create mode 100644 internal/store/store.go diff --git a/.env.example b/.env.example index adf6d4b..a12e74c 100644 --- a/.env.example +++ b/.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 diff --git a/cmd/gpt-load/main.go b/cmd/gpt-load/main.go index 45e1d4a..aca736c 100644 --- a/cmd/gpt-load/main.go +++ b/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) diff --git a/go.mod b/go.mod index 4c5a5e7..3ebe6e7 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index 33215af..acb8e11 100644 --- a/go.sum +++ b/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= diff --git a/internal/config/manager.go b/internal/config/manager.go index 39b49a3..258d517 100644 --- a/internal/config/manager.go +++ b/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 diff --git a/internal/store/factory.go b/internal/store/factory.go new file mode 100644 index 0000000..63bdf66 --- /dev/null +++ b/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 +} diff --git a/internal/store/memory.go b/internal/store/memory.go new file mode 100644 index 0000000..6904042 --- /dev/null +++ b/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 +} diff --git a/internal/store/redis.go b/internal/store/redis.go new file mode 100644 index 0000000..634d796 --- /dev/null +++ b/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() +} diff --git a/internal/store/store.go b/internal/store/store.go new file mode 100644 index 0000000..f4a1fac --- /dev/null +++ b/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 +} diff --git a/internal/types/types.go b/internal/types/types.go index abf5da8..17384c2 100644 --- a/internal/types/types.go +++ b/internal/types/types.go @@ -12,6 +12,7 @@ type ConfigManager interface { GetPerformanceConfig() PerformanceConfig GetLogConfig() LogConfig GetEffectiveServerConfig() ServerConfig + GetRedisDSN() string Validate() error DisplayConfig() ReloadConfig() error