feat(go-redis): store surrogate tags as native sets

core/registered.go

@@ -3,8 +3,42 @@ package core
 import (
 	"fmt"
 	"sync"
+	"time"
+
+	"github.com/pierrec/lz4/v4"
 )
 
+// Lz4WriterPool pools lz4 writers, which are safe to reuse once Close has
+// flushed the frame. Callers must Reset the writer before use and only
+// return it to the pool after Close. Readers must never be pooled this way:
+// a pooled reader escapes through http.Response.Body and would be recycled
+// while another goroutine still reads from it.
+var Lz4WriterPool = sync.Pool{New: func() any { return lz4.NewWriter(nil) }}
+
+// MappingWalker is an optional interface a Storer can implement to stream
+// mapping entries in bounded batches instead of materializing the whole
+// mapping index in memory like MapKeys does. The walk stops early when fn
+// returns false. The key passed to fn is stripped of the given prefix.
+type MappingWalker interface {
+	WalkMappings(prefix string, fn func(key string, value []byte) bool) error
+}
+
+// SetStorer is an optional interface a Storer can implement to store a set
+// of members under a key natively, instead of a separator-joined string
+// value that must be fully read and rewritten on every addition.
+type SetStorer interface {
+	// AddToSet adds members to the set stored at key, extending the set
+	// lifetime to at least the given duration when it is positive, without
+	// shortening a longer remaining lifetime.
+	AddToSet(key string, members []string, duration time.Duration) error
+	// GetSet returns all members of the set stored at key.
+	GetSet(key string) []string
+	// WalkSets visits every set whose key matches the prefix. The walk stops
+	// early when fn returns false. The key passed to fn is stripped of the
+	// given prefix.
+	WalkSets(prefix string, fn func(key string, members []string) bool) error
+}
+
 var registered = sync.Map{}
 
 func RegisterStorage(s Storer) {

go-redis/go-redis.go

@@ -155,31 +155,174 @@ func (provider *Redis) ListKeys() []string {
 // MapKeys method returns the list of existing keys.
 func (provider *Redis) MapKeys(prefix string) map[string]string {
 	mapKeys := map[string]string{}
-	keys := []string{}
 
-	iter := provider.inClient.Scan(provider.ctx, 0, prefix+"*", 100).Iterator()
+	_ = provider.WalkMappings(prefix, func(key string, value []byte) bool {
+		mapKeys[key] = string(value)
+
+		return true
+	})
+
+	return mapKeys
+}
+
+const mappingBatchSize = 100
+
+// WalkMappings streams the keys matching the prefix and their values in
+// bounded batches so the whole mapping index is never loaded in memory at
+// once. The walk stops early when walkFn returns false.
+func (provider *Redis) WalkMappings(prefix string, walkFn func(key string, value []byte) bool) error {
+	if provider.reconnecting {
+		provider.logger.Error("Impossible to walk the redis mappings while reconnecting.")
+
+		return errors.New("reconnecting error")
+	}
+
+	batch := make([]string, 0, mappingBatchSize)
+
+	flush := func() (bool, error) {
+		if len(batch) == 0 {
+			return true, nil
+		}
+
+		vals, err := provider.inClient.MGet(provider.ctx, batch...).Result()
+		if err != nil {
+			return false, err
+		}
+
+		for idx, item := range batch {
+			if idx >= len(vals) || vals[idx] == nil {
+				continue
+			}
+
+			value, ok := vals[idx].(string)
+			if !ok {
+				continue
+			}
+
+			k, _ := strings.CutPrefix(item, prefix)
+			if !walkFn(k, []byte(value)) {
+				return false, nil
+			}
+		}
+
+		batch = batch[:0]
+
+		return true, nil
+	}
+
+	iter := provider.inClient.Scan(provider.ctx, 0, prefix+"*", mappingBatchSize).Iterator()
 	for iter.Next(provider.ctx) {
-		keys = append(keys, iter.Val())
+		batch = append(batch, iter.Val())
+
+		if len(batch) >= mappingBatchSize {
+			if cont, err := flush(); err != nil || !cont {
+				return err
+			}
+		}
 	}
 
 	if err := iter.Err(); err != nil {
-		return mapKeys
+		return err
+	}
+
+	_, err := flush()
+
+	return err
+}
+
+// AddToSet stores members in the native set at key, migrating any legacy
+// string value first. A positive duration extends the set lifetime without
+// ever shortening a longer remaining one, and bounds legacy unbounded keys.
+func (provider *Redis) AddToSet(key string, members []string, duration time.Duration) error {
+	if provider.reconnecting {
+		provider.logger.Error("Impossible to add to the redis set while reconnecting.")
+
+		return errors.New("reconnecting error")
 	}
 
-	vals, err := provider.inClient.MGet(provider.ctx, keys...).Result()
+	legacy := provider.legacySetMembers(key)
+
+	values := make([]interface{}, 0, len(members)+len(legacy))
+	for _, member := range members {
+		values = append(values, member)
+	}
+
+	for _, member := range legacy {
+		values = append(values, member)
+	}
+
+	expire := time.Duration(0)
+
+	if duration > 0 {
+		if remaining := provider.inClient.TTL(provider.ctx, key).Val(); remaining < duration {
+			expire = duration
+		}
+	}
+
+	_, err := provider.inClient.TxPipelined(provider.ctx, func(pipe redis.Pipeliner) error {
+		if len(legacy) > 0 {
+			pipe.Del(provider.ctx, key)
+		}
+
+		pipe.SAdd(provider.ctx, key, values...)
+
+		if expire > 0 {
+			pipe.Expire(provider.ctx, key, expire)
+		}
+
+		return nil
+	})
+	if err != nil {
+		provider.logger.Errorf("Impossible to add members to the set %s into Redis, %v", key, err)
+	}
+
+	return err
+}
+
+// GetSet returns all members of the set stored at key, supporting sets that
+// are still stored in the legacy string format.
+func (provider *Redis) GetSet(key string) []string {
+	if provider.reconnecting {
+		provider.logger.Error("Impossible to get the redis set while reconnecting.")
+
+		return nil
+	}
+
+	if legacy := provider.legacySetMembers(key); len(legacy) > 0 {
+		return legacy
+	}
+
+	members, err := provider.inClient.SMembers(provider.ctx, key).Result()
 	if err != nil {
-		return mapKeys
+		return nil
+	}
+
+	return members
+}
+
+// WalkSets visits every set whose key matches the prefix. The walk stops
+// early when walkFn returns false.
+func (provider *Redis) WalkSets(prefix string, walkFn func(key string, members []string) bool) error {
+	if provider.reconnecting {
+		provider.logger.Error("Impossible to walk the redis sets while reconnecting.")
+
+		return errors.New("reconnecting error")
 	}
 
-	for idx, item := range keys {
-		k, _ := strings.CutPrefix(item, prefix)
+	iter := provider.inClient.Scan(provider.ctx, 0, prefix+"*", mappingBatchSize).Iterator()
+	for iter.Next(provider.ctx) {
+		members := provider.GetSet(iter.Val())
+		if len(members) == 0 {
+			continue
+		}
 
-		if vals[idx] != nil {
-			mapKeys[k] = vals[idx].(string)
+		key, _ := strings.CutPrefix(iter.Val(), prefix)
+		if !walkFn(key, members) {
+			return nil
 		}
 	}
 
-	return mapKeys
+	return iter.Err()
 }
 
 // GetMultiLevel tries to load the key and check if one of linked keys is a fresh/stale candidate.
@@ -201,6 +344,16 @@ func (provider *Redis) SetMultiLevel(baseKey, variedKey string, value []byte, va
 	compressed := new(bytes.Buffer)
 	writer := lz4.NewWriter(compressed)
 
+	// The lz4 default block size is 4 MB, which makes every compression and
+	// later decompression of the value churn 4 MB pooled blocks even for tiny
+	// payloads. Cached bodies are usually far smaller, so use the smallest
+	// block size. Readers pick the block size up from the frame header.
+	if err := writer.Apply(lz4.BlockSizeOption(lz4.Block64Kb)); err != nil {
+		provider.logger.Errorf("Impossible to configure the compressor for key %s into Redis, %v", variedKey, err)
+
+		return err
+	}
+
 	if _, err := writer.Write(value); err != nil {
 		_ = writer.Close()
 
@@ -233,7 +386,17 @@ func (provider *Redis) SetMultiLevel(baseKey, variedKey string, value []byte, va
 		return err
 	}
 
-	if err = provider.Set(mappingKey, val, -1); err != nil {
+	// Bound the mapping key lifetime instead of storing it forever: it only
+	// needs to outlive the longest-lived entry it references. Never shorten
+	// an expiration owned by a longer-lived entry; TTL returns a negative
+	// value for missing keys or keys without expiration, so legacy unbounded
+	// mapping keys become bounded on their next update.
+	mappingTTL := duration + provider.stale
+	if remaining := provider.inClient.TTL(provider.ctx, mappingKey).Val(); remaining > mappingTTL {
+		mappingTTL = remaining
+	}
+
+	if err = provider.inClient.Set(provider.ctx, mappingKey, val, mappingTTL).Err(); err != nil {
 		provider.logger.Errorf("Impossible to set value into Redis, %v", err)
 	}
 
@@ -370,3 +533,18 @@ func (provider *Redis) Reconnect() {
 		provider.Reconnect()
 	}
 }
+
+// legacySetMembers returns the members of a set that is still stored in the
+// legacy format: a single comma-joined string value.
+func (provider *Redis) legacySetMembers(key string) []string {
+	if keyType, _ := provider.inClient.Type(provider.ctx, key).Result(); keyType != "string" {
+		return nil
+	}
+
+	value, err := provider.inClient.Get(provider.ctx, key).Result()
+	if err != nil || value == "" {
+		return nil
+	}
+
+	return strings.Split(value, ",")
+}