feat: 增加内存存储

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