feat: 增加内存存储

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
+}