Fix basepool & 优化 repo

2025-11-23 22:42:58 +08:00
parent 2b0b9b67dc
commit 6c7283d51b
16 changed files with 1312 additions and 723 deletions
--- a/internal/store/memory_store.go
+++ b/internal/store/memory_store.go
@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"math/rand"
 	"sort"
+	"strconv"
 	"sync"
 	"time"

@@ -65,7 +66,7 @@ func (s *memoryStore) startGCollector() {
 	}
 }

-// --- [架构修正] 所有方法签名都增加了 context.Context 参数以匹配接口 ---
+// --- 所有方法签名都增加了 context.Context 参数以匹配接口 ---
 // --- 内存实现可以忽略该参数，用 _ 接收 ---

 func (s *memoryStore) Set(_ context.Context, key string, value []byte, ttl time.Duration) error {
@@ -108,6 +109,17 @@ func (s *memoryStore) Exists(_ context.Context, key string) (bool, error) {
 	return ok && !item.isExpired(), nil
 }

+func (s *memoryStore) Expire(_ context.Context, key string, expiration time.Duration) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	item, ok := s.items[key]
+	if !ok {
+		return ErrNotFound
+	}
+	item.expireAt = time.Now().Add(expiration)
+	return nil
+}
+
 func (s *memoryStore) SetNX(_ context.Context, key string, value []byte, ttl time.Duration) (bool, error) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
@@ -159,6 +171,21 @@ func (s *memoryStore) HSet(_ context.Context, key string, values map[string]any)
 	return nil
 }

+func (s *memoryStore) HGet(_ context.Context, key, field string) (string, error) {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+	item, ok := s.items[key]
+	if !ok || item.isExpired() {
+		return "", ErrNotFound
+	}
+	if hash, ok := item.value.(map[string]string); ok {
+		if value, exists := hash[field]; exists {
+			return value, nil
+		}
+	}
+	return "", ErrNotFound
+}
+
 func (s *memoryStore) HGetAll(_ context.Context, key string) (map[string]string, error) {
 	s.mu.RLock()
 	defer s.mu.RUnlock()
@@ -351,6 +378,26 @@ func (s *memoryStore) SRandMember(_ context.Context, key string) (string, error)
 	return members[n], nil
 }

+func (s *memoryStore) SUnionStore(_ context.Context, destination string, keys ...string) (int64, error) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	unionSet := make(map[string]struct{})
+	for _, key := range keys {
+		item, ok := s.items[key]
+		if !ok || item.isExpired() {
+			continue
+		}
+		if set, ok := item.value.(map[string]struct{}); ok {
+			for member := range set {
+				unionSet[member] = struct{}{}
+			}
+		}
+	}
+	destItem := &memoryStoreItem{value: unionSet}
+	s.items[destination] = destItem
+	return int64(len(unionSet)), nil
+}
+
 func (s *memoryStore) Rotate(_ context.Context, key string) (string, error) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
@@ -388,6 +435,16 @@ func (s *memoryStore) LIndex(_ context.Context, key string, index int64) (string
 	return list[index], nil
 }

+func (s *memoryStore) MSet(ctx context.Context, values map[string]any) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	for key, value := range values {
+		// 内存存储不支持独立的 TTL，因此我们假设永不过期
+		s.items[key] = &memoryStoreItem{value: value, expireAt: time.Time{}}
+	}
+	return nil
+}
+
 func (s *memoryStore) ZAdd(_ context.Context, key string, members map[string]float64) error {
 	s.mu.Lock()
 	defer s.mu.Unlock()
@@ -556,6 +613,22 @@ func (p *memoryPipeliner) Del(keys ...string) {
 		}
 	})
 }
+
+func (p *memoryPipeliner) Set(key string, value []byte, expiration time.Duration) {
+	capturedKey := key
+	capturedValue := value
+	p.ops = append(p.ops, func() {
+		var expireAt time.Time
+		if expiration > 0 {
+			expireAt = time.Now().Add(expiration)
+		}
+		p.store.items[capturedKey] = &memoryStoreItem{
+			value:    capturedValue,
+			expireAt: expireAt,
+		}
+	})
+}
+
 func (p *memoryPipeliner) SAdd(key string, members ...any) {
 	capturedKey := key
 	capturedMembers := make([]any, len(members))
@@ -576,6 +649,7 @@ func (p *memoryPipeliner) SAdd(key string, members ...any) {
 		}
 	})
 }
+
 func (p *memoryPipeliner) SRem(key string, members ...any) {
 	capturedKey := key
 	capturedMembers := make([]any, len(members))
@@ -615,11 +689,125 @@ func (p *memoryPipeliner) LPush(key string, values ...any) {
 		item.value = append(stringValues, list...)
 	})
 }
-func (p *memoryPipeliner) LRem(key string, count int64, value any)     {}
-func (p *memoryPipeliner) HSet(key string, values map[string]any)      {}
-func (p *memoryPipeliner) HIncrBy(key, field string, incr int64)       {}
-func (p *memoryPipeliner) ZAdd(key string, members map[string]float64) {}
-func (p *memoryPipeliner) ZRem(key string, members ...any)             {}
+func (p *memoryPipeliner) LRem(key string, count int64, value any) {
+	capturedKey := key
+	capturedValue := fmt.Sprintf("%v", value)
+	p.ops = append(p.ops, func() {
+		item, ok := p.store.items[capturedKey]
+		if !ok || item.isExpired() {
+			return
+		}
+		list, ok := item.value.([]string)
+		if !ok {
+			return
+		}
+		newList := make([]string, 0, len(list))
+		removed := int64(0)
+		for _, v := range list {
+			if count != 0 && v == capturedValue && (count < 0 || removed < count) {
+				removed++
+				continue
+			}
+			newList = append(newList, v)
+		}
+		item.value = newList
+	})
+}
+func (p *memoryPipeliner) HSet(key string, values map[string]any) {
+	capturedKey := key
+	capturedValues := make(map[string]any, len(values))
+	for k, v := range values {
+		capturedValues[k] = v
+	}
+	p.ops = append(p.ops, func() {
+		item, ok := p.store.items[capturedKey]
+		if !ok || item.isExpired() {
+			item = &memoryStoreItem{value: make(map[string]string)}
+			p.store.items[capturedKey] = item
+		}
+		hash, ok := item.value.(map[string]string)
+		if !ok {
+			hash = make(map[string]string)
+			item.value = hash
+		}
+		for field, value := range capturedValues {
+			hash[field] = fmt.Sprintf("%v", value)
+		}
+	})
+}
+func (p *memoryPipeliner) HIncrBy(key, field string, incr int64) {
+	capturedKey := key
+	capturedField := field
+	p.ops = append(p.ops, func() {
+		item, ok := p.store.items[capturedKey]
+		if !ok || item.isExpired() {
+			item = &memoryStoreItem{value: make(map[string]string)}
+			p.store.items[capturedKey] = item
+		}
+		hash, ok := item.value.(map[string]string)
+		if !ok {
+			hash = make(map[string]string)
+			item.value = hash
+		}
+		current, _ := strconv.ParseInt(hash[capturedField], 10, 64)
+		hash[capturedField] = strconv.FormatInt(current+incr, 10)
+	})
+}
+func (p *memoryPipeliner) ZAdd(key string, members map[string]float64) {
+	capturedKey := key
+	capturedMembers := make(map[string]float64, len(members))
+	for k, v := range members {
+		capturedMembers[k] = v
+	}
+	p.ops = append(p.ops, func() {
+		item, ok := p.store.items[capturedKey]
+		if !ok || item.isExpired() {
+			item = &memoryStoreItem{value: make(map[string]float64)}
+			p.store.items[capturedKey] = item
+		}
+		zset, ok := item.value.(map[string]float64)
+		if !ok {
+			zset = make(map[string]float64)
+			item.value = zset
+		}
+		for member, score := range capturedMembers {
+			zset[member] = score
+		}
+	})
+}
+func (p *memoryPipeliner) ZRem(key string, members ...any) {
+	capturedKey := key
+	capturedMembers := make([]any, len(members))
+	copy(capturedMembers, members)
+	p.ops = append(p.ops, func() {
+		item, ok := p.store.items[capturedKey]
+		if !ok || item.isExpired() {
+			return
+		}
+		zset, ok := item.value.(map[string]float64)
+		if !ok {
+			return
+		}
+		for _, member := range capturedMembers {
+			delete(zset, fmt.Sprintf("%v", member))
+		}
+	})
+}
+
+func (p *memoryPipeliner) MSet(values map[string]any) {
+	capturedValues := make(map[string]any, len(values))
+	for k, v := range values {
+		capturedValues[k] = v
+	}
+	p.ops = append(p.ops, func() {
+		for key, value := range capturedValues {
+			p.store.items[key] = &memoryStoreItem{
+				value:    value,
+				expireAt: time.Time{}, // Pipelined MSet 同样假设永不过期
+			}
+		}
+	})
+}

 type memorySubscription struct {
 	store       *memoryStore