feat(redis): add shard-aware key schema, Lua scripts, and atomic ops
Implements the kmsvc: key schema (inflight, pending/watermark keyed by shard+partition, vis_index, dedup, fifo_lock, queue meta, shard map) plus the atomic reap/ack Lua scripts used for safe multi-replica redelivery and DLQ routing.
This commit is contained in:
riotpiaole
@@ -0,0 +1,12 @@
|
||||
package redis
|
||||
|
||||
import "github.com/redis/go-redis/v9"
|
||||
|
||||
// NewClient builds a Redis client for the given address/password/db.
|
||||
func NewClient(addr, password string, db int) *redis.Client {
|
||||
return redis.NewClient(&redis.Options{
|
||||
Addr: addr,
|
||||
Password: password,
|
||||
DB: db,
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,21 @@
|
||||
package redis
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
// TryDedup attempts to claim a MessageDeduplicationId for a FIFO MessageGroupId
|
||||
// (design.md §4 `kmsvc:dedup:` row). Returns true if this is the first send
|
||||
// within the dedup window (the message should be accepted), false if it's a
|
||||
// duplicate (the message should be silently deduped).
|
||||
func TryDedup(ctx context.Context, rdb *redis.Client, queue, groupID, dedupID string, window time.Duration) (bool, error) {
|
||||
ok, err := rdb.SetNX(ctx, DedupKey(queue, groupID, dedupID), "1", window).Result()
|
||||
if err != nil {
|
||||
return false, fmt.Errorf("dedup check %s/%s: %w", groupID, dedupID, err)
|
||||
}
|
||||
return ok, nil
|
||||
}
|
||||
@@ -0,0 +1,92 @@
|
||||
package redis
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
// InFlightRecord is the cached state of a message currently checked out by a
|
||||
// consumer, per design.md §4's `kmsvc:inflight:` row.
|
||||
type InFlightRecord struct {
|
||||
ShardID string
|
||||
Topic string
|
||||
Partition int32
|
||||
Offset int64
|
||||
GroupID string
|
||||
DedupID string
|
||||
ReceiveCount int32
|
||||
Body string
|
||||
}
|
||||
|
||||
// PutInFlight records a freshly received message as in-flight: writes the
|
||||
// hash, sets its TTL, and adds it to the visibility-expiry index so the
|
||||
// reaper (§5) can find it once visibleAt passes.
|
||||
func PutInFlight(ctx context.Context, rdb *redis.Client, queue, receiptHandle string, rec InFlightRecord, visibilityTimeout time.Duration) error {
|
||||
key := InFlightKey(queue, receiptHandle)
|
||||
visibleAt := time.Now().Add(visibilityTimeout)
|
||||
|
||||
pipe := rdb.TxPipeline()
|
||||
pipe.HSet(ctx, key, map[string]any{
|
||||
"shardId": rec.ShardID,
|
||||
"topic": rec.Topic,
|
||||
"partition": rec.Partition,
|
||||
"offset": rec.Offset,
|
||||
"groupId": rec.GroupID,
|
||||
"dedupId": rec.DedupID,
|
||||
"receiveCount": rec.ReceiveCount,
|
||||
"body": rec.Body,
|
||||
})
|
||||
pipe.Expire(ctx, key, visibilityTimeout+time.Minute)
|
||||
pipe.ZAdd(ctx, VisIndexKey(queue), redis.Z{Score: float64(visibleAt.UnixMilli()), Member: receiptHandle})
|
||||
_, err := pipe.Exec(ctx)
|
||||
if err != nil {
|
||||
return fmt.Errorf("put in-flight %s: %w", receiptHandle, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetInFlight reads a message's in-flight record, or ok=false if it's gone
|
||||
// (already acked or already DLQ-routed).
|
||||
func GetInFlight(ctx context.Context, rdb *redis.Client, queue, receiptHandle string) (InFlightRecord, bool, error) {
|
||||
m, err := rdb.HGetAll(ctx, InFlightKey(queue, receiptHandle)).Result()
|
||||
if err != nil {
|
||||
return InFlightRecord{}, false, fmt.Errorf("get in-flight %s: %w", receiptHandle, err)
|
||||
}
|
||||
if len(m) == 0 {
|
||||
return InFlightRecord{}, false, nil
|
||||
}
|
||||
partition, _ := strconv.ParseInt(m["partition"], 10, 32)
|
||||
offset, _ := strconv.ParseInt(m["offset"], 10, 64)
|
||||
receiveCount, _ := strconv.ParseInt(m["receiveCount"], 10, 32)
|
||||
return InFlightRecord{
|
||||
ShardID: m["shardId"],
|
||||
Topic: m["topic"],
|
||||
Partition: int32(partition),
|
||||
Offset: offset,
|
||||
GroupID: m["groupId"],
|
||||
DedupID: m["dedupId"],
|
||||
ReceiveCount: int32(receiveCount),
|
||||
Body: m["body"],
|
||||
}, true, nil
|
||||
}
|
||||
|
||||
// ExtendVisibility implements ChangeMessageVisibility: bumps the vis_index
|
||||
// score for an in-flight message without touching its receive count.
|
||||
func ExtendVisibility(ctx context.Context, rdb *redis.Client, queue, receiptHandle string, newTimeout time.Duration) (bool, error) {
|
||||
exists, err := rdb.Exists(ctx, InFlightKey(queue, receiptHandle)).Result()
|
||||
if err != nil {
|
||||
return false, fmt.Errorf("extend visibility %s: %w", receiptHandle, err)
|
||||
}
|
||||
if exists == 0 {
|
||||
return false, nil
|
||||
}
|
||||
visibleAt := time.Now().Add(newTimeout)
|
||||
if err := rdb.ZAdd(ctx, VisIndexKey(queue), redis.Z{Score: float64(visibleAt.UnixMilli()), Member: receiptHandle}).Err(); err != nil {
|
||||
return false, fmt.Errorf("extend visibility %s: %w", receiptHandle, err)
|
||||
}
|
||||
return true, nil
|
||||
}
|
||||
@@ -0,0 +1,53 @@
|
||||
// Package redis implements the key schemas and atomic Lua scripts from
|
||||
// design.md §4/§5, shared by the queue-operator and the message-plane service.
|
||||
package redis
|
||||
|
||||
import "fmt"
|
||||
|
||||
func InFlightKey(queue, receiptHandle string) string {
|
||||
return fmt.Sprintf("kmsvc:inflight:%s:%s", queue, receiptHandle)
|
||||
}
|
||||
|
||||
func PendingKey(queue, shardID string, partition int32) string {
|
||||
return fmt.Sprintf("kmsvc:pending:%s:%s:%d", queue, shardID, partition)
|
||||
}
|
||||
|
||||
func WatermarkKey(queue, shardID string, partition int32) string {
|
||||
return fmt.Sprintf("kmsvc:watermark:%s:%s:%d", queue, shardID, partition)
|
||||
}
|
||||
|
||||
func VisIndexKey(queue string) string {
|
||||
return fmt.Sprintf("kmsvc:vis_index:%s", queue)
|
||||
}
|
||||
|
||||
func DedupKey(queue, groupID, dedupID string) string {
|
||||
return fmt.Sprintf("kmsvc:dedup:%s:%s:%s", queue, groupID, dedupID)
|
||||
}
|
||||
|
||||
func FIFOLockKey(queue, groupID string) string {
|
||||
return fmt.Sprintf("kmsvc:fifo_lock:%s:%s", queue, groupID)
|
||||
}
|
||||
|
||||
func QueueMetaKey(queue string) string {
|
||||
return fmt.Sprintf("kmsvc:queue:%s", queue)
|
||||
}
|
||||
|
||||
func ShardMapKey(queue string) string {
|
||||
return fmt.Sprintf("kmsvc:shardmap:%s", queue)
|
||||
}
|
||||
|
||||
func ShardMapChannel(queue string) string {
|
||||
return fmt.Sprintf("kmsvc:shardmap_invalidate:%s", queue)
|
||||
}
|
||||
|
||||
func QueueMetaChannel(queue string) string {
|
||||
return fmt.Sprintf("kmsvc:queuemeta_invalidate:%s", queue)
|
||||
}
|
||||
|
||||
func RedeliverKey(queue string) string {
|
||||
return fmt.Sprintf("kmsvc:redeliver:%s", queue)
|
||||
}
|
||||
|
||||
func QueueLockKey(queue string) string {
|
||||
return fmt.Sprintf("kmsvc:queue_lock:%s", queue)
|
||||
}
|
||||
@@ -0,0 +1,29 @@
|
||||
package redis
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestKeyFormats(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
got string
|
||||
want string
|
||||
}{
|
||||
{"inflight", InFlightKey("orders", "rh-1"), "kmsvc:inflight:orders:rh-1"},
|
||||
{"pending", PendingKey("orders", "0", 3), "kmsvc:pending:orders:0:3"},
|
||||
{"watermark", WatermarkKey("orders", "0", 3), "kmsvc:watermark:orders:0:3"},
|
||||
{"vis_index", VisIndexKey("orders"), "kmsvc:vis_index:orders"},
|
||||
{"dedup", DedupKey("orders", "g1", "d1"), "kmsvc:dedup:orders:g1:d1"},
|
||||
{"fifo_lock", FIFOLockKey("orders", "g1"), "kmsvc:fifo_lock:orders:g1"},
|
||||
{"queue", QueueMetaKey("orders"), "kmsvc:queue:orders"},
|
||||
{"shardmap", ShardMapKey("orders"), "kmsvc:shardmap:orders"},
|
||||
{"redeliver", RedeliverKey("orders"), "kmsvc:redeliver:orders"},
|
||||
{"queue_lock", QueueLockKey("orders"), "kmsvc:queue_lock:orders"},
|
||||
}
|
||||
for _, c := range cases {
|
||||
t.Run(c.name, func(t *testing.T) {
|
||||
if c.got != c.want {
|
||||
t.Errorf("got %q, want %q", c.got, c.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,37 @@
|
||||
-- Atomic ack for DeleteMessage (design.md §3).
|
||||
-- KEYS[1] = inflight key
|
||||
-- KEYS[2] = vis_index key (queue-level)
|
||||
-- ARGV[1] = receiptHandle
|
||||
-- ARGV[2] = queueName
|
||||
--
|
||||
-- Returns {"not_found"} if the message is already gone (already acked, or
|
||||
-- already DLQ-routed/redelivered by the reaper).
|
||||
-- Returns {"acked"} otherwise.
|
||||
|
||||
local inflight = redis.call('HGETALL', KEYS[1])
|
||||
if #inflight == 0 then
|
||||
return {'not_found'}
|
||||
end
|
||||
|
||||
local m = {}
|
||||
for i = 1, #inflight, 2 do
|
||||
m[inflight[i]] = inflight[i + 1]
|
||||
end
|
||||
|
||||
local queueName = ARGV[2]
|
||||
|
||||
redis.call('ZREM', KEYS[2], ARGV[1])
|
||||
|
||||
local pendingKey = 'kmsvc:pending:' .. queueName .. ':' .. m['shardId'] .. ':' .. m['partition']
|
||||
redis.call('ZREM', pendingKey, m['offset'])
|
||||
redis.call('DEL', KEYS[1])
|
||||
|
||||
if m['groupId'] and m['groupId'] ~= '' then
|
||||
local lockKey = 'kmsvc:fifo_lock:' .. queueName .. ':' .. m['groupId']
|
||||
local lockVal = redis.call('GET', lockKey)
|
||||
if lockVal == ARGV[1] then
|
||||
redis.call('DEL', lockKey)
|
||||
end
|
||||
end
|
||||
|
||||
return {'acked'}
|
||||
@@ -0,0 +1,49 @@
|
||||
-- Atomic check-and-act for one expired in-flight message (design.md §5).
|
||||
-- KEYS[1] = inflight key
|
||||
-- KEYS[2] = vis_index key (queue-level)
|
||||
-- ARGV[1] = receiptHandle
|
||||
-- ARGV[2] = maxReceiveCount
|
||||
-- ARGV[3] = queueName
|
||||
--
|
||||
-- Returns {"gone"} if another reaper already won the race for this key.
|
||||
-- Returns {"dlq", topic, shardId, partition, offset, body, groupId, dedupId}
|
||||
-- when receiveCount has reached maxReceiveCount.
|
||||
-- Returns {"redeliver", receiptHandle, newReceiveCount} otherwise.
|
||||
|
||||
local removed = redis.call('ZREM', KEYS[2], ARGV[1])
|
||||
if removed == 0 then
|
||||
return {'gone'}
|
||||
end
|
||||
|
||||
local inflight = redis.call('HGETALL', KEYS[1])
|
||||
if #inflight == 0 then
|
||||
return {'gone'}
|
||||
end
|
||||
|
||||
local m = {}
|
||||
for i = 1, #inflight, 2 do
|
||||
m[inflight[i]] = inflight[i + 1]
|
||||
end
|
||||
|
||||
local queueName = ARGV[3]
|
||||
local receiveCount = tonumber(m['receiveCount'] or '0')
|
||||
local maxReceive = tonumber(ARGV[2])
|
||||
|
||||
if m['groupId'] and m['groupId'] ~= '' then
|
||||
local lockKey = 'kmsvc:fifo_lock:' .. queueName .. ':' .. m['groupId']
|
||||
local lockVal = redis.call('GET', lockKey)
|
||||
if lockVal == ARGV[1] then
|
||||
redis.call('DEL', lockKey)
|
||||
end
|
||||
end
|
||||
|
||||
if receiveCount >= maxReceive then
|
||||
local pendingKey = 'kmsvc:pending:' .. queueName .. ':' .. m['shardId'] .. ':' .. m['partition']
|
||||
redis.call('ZREM', pendingKey, m['offset'])
|
||||
redis.call('DEL', KEYS[1])
|
||||
return {'dlq', m['topic'] or '', m['shardId'] or '', m['partition'] or '', m['offset'] or '', m['body'] or '', m['groupId'] or '', m['dedupId'] or ''}
|
||||
end
|
||||
|
||||
redis.call('HSET', KEYS[1], 'receiveCount', receiveCount + 1)
|
||||
redis.call('RPUSH', 'kmsvc:redeliver:' .. queueName, ARGV[1])
|
||||
return {'redeliver', ARGV[1], tostring(receiveCount + 1)}
|
||||
@@ -0,0 +1,81 @@
|
||||
package redis
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
// QueueMeta is the queue-operator-written config snapshot a message-plane
|
||||
// replica reads to know how to serve a queue, per design.md §4's `kmsvc:queue:` row.
|
||||
type QueueMeta struct {
|
||||
FIFO bool
|
||||
VisibilityTimeoutSeconds int32
|
||||
MaxReceiveCount int32
|
||||
DLQQueueName string
|
||||
PartitionsPerShard int32
|
||||
RetentionSeconds int32
|
||||
CreatedAt time.Time
|
||||
}
|
||||
|
||||
// PutQueueMeta writes a queue's config and publishes an invalidation so any
|
||||
// in-process caches refresh on next read.
|
||||
func PutQueueMeta(ctx context.Context, rdb *redis.Client, queue string, m QueueMeta) error {
|
||||
key := QueueMetaKey(queue)
|
||||
createdAt := m.CreatedAt
|
||||
if createdAt.IsZero() {
|
||||
createdAt = time.Now()
|
||||
}
|
||||
pipe := rdb.TxPipeline()
|
||||
pipe.HSet(ctx, key, map[string]any{
|
||||
"fifo": m.FIFO,
|
||||
"visibilityTimeoutSeconds": m.VisibilityTimeoutSeconds,
|
||||
"maxReceiveCount": m.MaxReceiveCount,
|
||||
"dlqQueueName": m.DLQQueueName,
|
||||
"partitionsPerShard": m.PartitionsPerShard,
|
||||
"retentionSeconds": m.RetentionSeconds,
|
||||
"createdAt": createdAt.Unix(),
|
||||
})
|
||||
pipe.Publish(ctx, QueueMetaChannel(queue), "invalidate")
|
||||
if _, err := pipe.Exec(ctx); err != nil {
|
||||
return fmt.Errorf("put queue meta %s: %w", queue, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetQueueMeta reads a queue's config, or ok=false if it doesn't exist.
|
||||
func GetQueueMeta(ctx context.Context, rdb *redis.Client, queue string) (QueueMeta, bool, error) {
|
||||
m, err := rdb.HGetAll(ctx, QueueMetaKey(queue)).Result()
|
||||
if err != nil {
|
||||
return QueueMeta{}, false, fmt.Errorf("get queue meta %s: %w", queue, err)
|
||||
}
|
||||
if len(m) == 0 {
|
||||
return QueueMeta{}, false, nil
|
||||
}
|
||||
visTimeout, _ := strconv.ParseInt(m["visibilityTimeoutSeconds"], 10, 32)
|
||||
maxReceive, _ := strconv.ParseInt(m["maxReceiveCount"], 10, 32)
|
||||
partitionsPerShard, _ := strconv.ParseInt(m["partitionsPerShard"], 10, 32)
|
||||
retention, _ := strconv.ParseInt(m["retentionSeconds"], 10, 32)
|
||||
createdAtUnix, _ := strconv.ParseInt(m["createdAt"], 10, 64)
|
||||
return QueueMeta{
|
||||
FIFO: m["fifo"] == "1",
|
||||
VisibilityTimeoutSeconds: int32(visTimeout),
|
||||
MaxReceiveCount: int32(maxReceive),
|
||||
DLQQueueName: m["dlqQueueName"],
|
||||
PartitionsPerShard: int32(partitionsPerShard),
|
||||
RetentionSeconds: int32(retention),
|
||||
CreatedAt: time.Unix(createdAtUnix, 0),
|
||||
}, true, nil
|
||||
}
|
||||
|
||||
// DeleteQueueMeta removes a queue's config, used by the queue-operator on
|
||||
// Queue deletion.
|
||||
func DeleteQueueMeta(ctx context.Context, rdb *redis.Client, queue string) error {
|
||||
if err := rdb.Del(ctx, QueueMetaKey(queue)).Err(); err != nil {
|
||||
return fmt.Errorf("delete queue meta %s: %w", queue, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,141 @@
|
||||
package redis
|
||||
|
||||
import (
|
||||
"context"
|
||||
_ "embed"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
//go:embed lua/reap.lua
|
||||
var reapScript string
|
||||
|
||||
//go:embed lua/ack.lua
|
||||
var ackScript string
|
||||
|
||||
var (
|
||||
reapSHA = redis.NewScript(reapScript)
|
||||
ackSHA = redis.NewScript(ackScript)
|
||||
)
|
||||
|
||||
// ReapOutcome is the action taken by reap.lua for one expired in-flight entry.
|
||||
type ReapOutcome string
|
||||
|
||||
const (
|
||||
ReapOutcomeGone ReapOutcome = "gone"
|
||||
ReapOutcomeDLQ ReapOutcome = "dlq"
|
||||
ReapOutcomeRedeliver ReapOutcome = "redeliver"
|
||||
)
|
||||
|
||||
// ReapResult carries enough of the in-flight record back to the caller to
|
||||
// act on it (produce to DLQ topic, or nothing further for a redeliver).
|
||||
type ReapResult struct {
|
||||
Outcome ReapOutcome
|
||||
Topic string
|
||||
ShardID string
|
||||
Partition int32
|
||||
Offset int64
|
||||
Body string
|
||||
GroupID string
|
||||
DedupID string
|
||||
ReceiveCount int32
|
||||
}
|
||||
|
||||
// Reap runs the atomic check-and-act script for one expired vis_index entry
|
||||
// (design.md §5). Safe to call concurrently from multiple reaper replicas —
|
||||
// only the first caller for a given receiptHandle gets a non-"gone" outcome.
|
||||
func Reap(ctx context.Context, rdb *redis.Client, queue, receiptHandle string, maxReceiveCount int32) (ReapResult, error) {
|
||||
keys := []string{InFlightKey(queue, receiptHandle), VisIndexKey(queue)}
|
||||
res, err := reapSHA.Run(ctx, rdb, keys, receiptHandle, maxReceiveCount, queue).StringSlice()
|
||||
if err != nil {
|
||||
return ReapResult{}, fmt.Errorf("reap %s: %w", receiptHandle, err)
|
||||
}
|
||||
if len(res) == 0 || res[0] == string(ReapOutcomeGone) {
|
||||
return ReapResult{Outcome: ReapOutcomeGone}, nil
|
||||
}
|
||||
|
||||
switch ReapOutcome(res[0]) {
|
||||
case ReapOutcomeDLQ:
|
||||
partition, _ := strconv.ParseInt(res[3], 10, 32)
|
||||
offset, _ := strconv.ParseInt(res[4], 10, 64)
|
||||
return ReapResult{
|
||||
Outcome: ReapOutcomeDLQ,
|
||||
Topic: res[1],
|
||||
ShardID: res[2],
|
||||
Partition: int32(partition),
|
||||
Offset: offset,
|
||||
Body: res[5],
|
||||
GroupID: res[6],
|
||||
DedupID: res[7],
|
||||
}, nil
|
||||
case ReapOutcomeRedeliver:
|
||||
receiveCount, _ := strconv.ParseInt(res[2], 10, 32)
|
||||
return ReapResult{Outcome: ReapOutcomeRedeliver, ReceiveCount: int32(receiveCount)}, nil
|
||||
default:
|
||||
return ReapResult{}, fmt.Errorf("reap %s: unexpected outcome %q", receiptHandle, res[0])
|
||||
}
|
||||
}
|
||||
|
||||
// AckOutcome is the result of running ack.lua.
|
||||
type AckOutcome string
|
||||
|
||||
const (
|
||||
AckOutcomeAcked AckOutcome = "acked"
|
||||
AckOutcomeNotFound AckOutcome = "not_found"
|
||||
)
|
||||
|
||||
// Ack runs the atomic ack script for DeleteMessage (design.md §3).
|
||||
func Ack(ctx context.Context, rdb *redis.Client, queue, receiptHandle string) (AckOutcome, error) {
|
||||
keys := []string{InFlightKey(queue, receiptHandle), VisIndexKey(queue)}
|
||||
res, err := ackSHA.Run(ctx, rdb, keys, receiptHandle, queue).StringSlice()
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("ack %s: %w", receiptHandle, err)
|
||||
}
|
||||
if len(res) == 0 {
|
||||
return AckOutcomeNotFound, nil
|
||||
}
|
||||
return AckOutcome(res[0]), nil
|
||||
}
|
||||
|
||||
// ExpiredVisIndexEntries returns up to limit receiptHandles whose visibility
|
||||
// has already passed (design.md §5's sweep: ZRANGEBYSCORE ... 0 now LIMIT 0
|
||||
// limit), oldest-expiry first.
|
||||
func ExpiredVisIndexEntries(ctx context.Context, rdb *redis.Client, queue string, limit int64) ([]string, error) {
|
||||
handles, err := rdb.ZRangeByScore(ctx, VisIndexKey(queue), &redis.ZRangeBy{
|
||||
Min: "0",
|
||||
Max: strconv.FormatInt(time.Now().UnixMilli(), 10),
|
||||
Count: limit,
|
||||
}).Result()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("expired vis index %s: %w", queue, err)
|
||||
}
|
||||
return handles, nil
|
||||
}
|
||||
|
||||
// PopRedeliverable pops the next receiptHandle queued for redelivery
|
||||
// (design.md §4's `kmsvc:redeliver:` row), or ok=false if none is pending.
|
||||
func PopRedeliverable(ctx context.Context, rdb *redis.Client, queue string) (receiptHandle string, ok bool, err error) {
|
||||
v, err := rdb.LPop(ctx, RedeliverKey(queue)).Result()
|
||||
if err == redis.Nil {
|
||||
return "", false, nil
|
||||
}
|
||||
if err != nil {
|
||||
return "", false, fmt.Errorf("pop redeliverable %s: %w", queue, err)
|
||||
}
|
||||
return v, true, nil
|
||||
}
|
||||
|
||||
// PushRedeliverable re-queues a receiptHandle for redelivery. Used by
|
||||
// ReceiveMessage when a popped (or freshly fetched) message can't be handed
|
||||
// out yet because its FIFO group still has another message in flight — it
|
||||
// goes back on the list instead of into vis_index, so a later poll retries
|
||||
// it without re-fetching from Kafka.
|
||||
func PushRedeliverable(ctx context.Context, rdb *redis.Client, queue, receiptHandle string) error {
|
||||
if err := rdb.RPush(ctx, RedeliverKey(queue), receiptHandle).Err(); err != nil {
|
||||
return fmt.Errorf("push redeliverable %s: %w", queue, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,285 @@
|
||||
package redis
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/alicebob/miniredis/v2"
|
||||
goredis "github.com/redis/go-redis/v9"
|
||||
|
||||
"github.com/rockliang/kafka-management-service/internal/kafka"
|
||||
)
|
||||
|
||||
func newTestClient(t *testing.T) *goredis.Client {
|
||||
t.Helper()
|
||||
mr, err := miniredis.Run()
|
||||
if err != nil {
|
||||
t.Fatalf("starting miniredis: %v", err)
|
||||
}
|
||||
t.Cleanup(mr.Close)
|
||||
return goredis.NewClient(&goredis.Options{Addr: mr.Addr()})
|
||||
}
|
||||
|
||||
func TestDedupFirstAcceptsSecondRejects(t *testing.T) {
|
||||
rdb := newTestClient(t)
|
||||
ctx := context.Background()
|
||||
|
||||
first, err := TryDedup(ctx, rdb, "orders", "g1", "d1", 5*time.Minute)
|
||||
if err != nil || !first {
|
||||
t.Fatalf("first dedup attempt: ok=%v err=%v, want ok=true", first, err)
|
||||
}
|
||||
second, err := TryDedup(ctx, rdb, "orders", "g1", "d1", 5*time.Minute)
|
||||
if err != nil || second {
|
||||
t.Fatalf("second dedup attempt: ok=%v err=%v, want ok=false", second, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestInFlightPutGetExtend(t *testing.T) {
|
||||
rdb := newTestClient(t)
|
||||
ctx := context.Background()
|
||||
|
||||
rec := InFlightRecord{ShardID: "0", Topic: "kmsvc.orders.shard-0", Partition: 2, Offset: 42, GroupID: "g1", Body: "hello"}
|
||||
if err := PutInFlight(ctx, rdb, "orders", "rh-1", rec, 30*time.Second); err != nil {
|
||||
t.Fatalf("PutInFlight: %v", err)
|
||||
}
|
||||
|
||||
got, ok, err := GetInFlight(ctx, rdb, "orders", "rh-1")
|
||||
if err != nil || !ok {
|
||||
t.Fatalf("GetInFlight: ok=%v err=%v", ok, err)
|
||||
}
|
||||
if got.Topic != rec.Topic || got.Offset != rec.Offset || got.Body != rec.Body {
|
||||
t.Errorf("GetInFlight = %+v, want %+v", got, rec)
|
||||
}
|
||||
|
||||
extended, err := ExtendVisibility(ctx, rdb, "orders", "rh-1", time.Minute)
|
||||
if err != nil || !extended {
|
||||
t.Fatalf("ExtendVisibility: ok=%v err=%v", extended, err)
|
||||
}
|
||||
|
||||
missing, err := ExtendVisibility(ctx, rdb, "orders", "rh-missing", time.Minute)
|
||||
if err != nil || missing {
|
||||
t.Fatalf("ExtendVisibility(missing): ok=%v err=%v, want ok=false", missing, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWatermarkAdvancesAsPendingDrains(t *testing.T) {
|
||||
rdb := newTestClient(t)
|
||||
ctx := context.Background()
|
||||
|
||||
for _, off := range []int64{10, 11, 12} {
|
||||
if err := AddPending(ctx, rdb, "orders", "0", 2, off); err != nil {
|
||||
t.Fatalf("AddPending(%d): %v", off, err)
|
||||
}
|
||||
}
|
||||
|
||||
min, ok, err := MinPending(ctx, rdb, "orders", "0", 2)
|
||||
if err != nil || !ok || min != 10 {
|
||||
t.Fatalf("MinPending = %d, ok=%v err=%v, want 10", min, ok, err)
|
||||
}
|
||||
|
||||
if err := rdb.ZRem(ctx, PendingKey("orders", "0", 2), 10).Err(); err != nil {
|
||||
t.Fatalf("ack offset 10: %v", err)
|
||||
}
|
||||
min, ok, err = MinPending(ctx, rdb, "orders", "0", 2)
|
||||
if err != nil || !ok || min != 11 {
|
||||
t.Fatalf("MinPending after ack = %d, ok=%v err=%v, want 11", min, ok, err)
|
||||
}
|
||||
|
||||
if err := SetWatermark(ctx, rdb, "orders", "0", 2, min-1); err != nil {
|
||||
t.Fatalf("SetWatermark: %v", err)
|
||||
}
|
||||
got, ok, err := GetWatermark(ctx, rdb, "orders", "0", 2)
|
||||
if err != nil || !ok || got != 10 {
|
||||
t.Fatalf("GetWatermark = %d, ok=%v err=%v, want 10", got, ok, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestQueueMetaRoundTrip(t *testing.T) {
|
||||
rdb := newTestClient(t)
|
||||
ctx := context.Background()
|
||||
|
||||
m := QueueMeta{FIFO: true, VisibilityTimeoutSeconds: 30, MaxReceiveCount: 5, PartitionsPerShard: 6, RetentionSeconds: 345600}
|
||||
if err := PutQueueMeta(ctx, rdb, "orders-fifo", m); err != nil {
|
||||
t.Fatalf("PutQueueMeta: %v", err)
|
||||
}
|
||||
got, ok, err := GetQueueMeta(ctx, rdb, "orders-fifo")
|
||||
if err != nil || !ok {
|
||||
t.Fatalf("GetQueueMeta: ok=%v err=%v", ok, err)
|
||||
}
|
||||
if got.FIFO != true || got.VisibilityTimeoutSeconds != 30 || got.MaxReceiveCount != 5 {
|
||||
t.Errorf("GetQueueMeta = %+v", got)
|
||||
}
|
||||
|
||||
if err := DeleteQueueMeta(ctx, rdb, "orders-fifo"); err != nil {
|
||||
t.Fatalf("DeleteQueueMeta: %v", err)
|
||||
}
|
||||
_, ok, err = GetQueueMeta(ctx, rdb, "orders-fifo")
|
||||
if err != nil || ok {
|
||||
t.Fatalf("GetQueueMeta after delete: ok=%v err=%v, want false", ok, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestShardMapRoundTrip(t *testing.T) {
|
||||
rdb := newTestClient(t)
|
||||
ctx := context.Background()
|
||||
|
||||
shards := []kafka.Shard{
|
||||
{ID: "0", Topic: "kmsvc.orders.fifo.shard-0", HashRangeStart: 0, HashRangeEnd: kafka.FullHashRangeEnd},
|
||||
}
|
||||
if err := PutShardMap(ctx, rdb, "orders-fifo", shards); err != nil {
|
||||
t.Fatalf("PutShardMap: %v", err)
|
||||
}
|
||||
got, ok, err := GetShardMap(ctx, rdb, "orders-fifo")
|
||||
if err != nil || !ok {
|
||||
t.Fatalf("GetShardMap: ok=%v err=%v", ok, err)
|
||||
}
|
||||
if len(got) != 1 || got[0].Topic != shards[0].Topic {
|
||||
t.Errorf("GetShardMap = %+v, want %+v", got, shards)
|
||||
}
|
||||
|
||||
if err := DeleteShardMap(ctx, rdb, "orders-fifo"); err != nil {
|
||||
t.Fatalf("DeleteShardMap: %v", err)
|
||||
}
|
||||
_, ok, err = GetShardMap(ctx, rdb, "orders-fifo")
|
||||
if err != nil || ok {
|
||||
t.Fatalf("GetShardMap after delete: ok=%v err=%v, want false", ok, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAckRemovesInFlightAndPending(t *testing.T) {
|
||||
rdb := newTestClient(t)
|
||||
ctx := context.Background()
|
||||
|
||||
rec := InFlightRecord{ShardID: "0", Topic: "kmsvc.orders.shard-0", Partition: 1, Offset: 7, GroupID: "g1"}
|
||||
if err := PutInFlight(ctx, rdb, "orders", "rh-1", rec, 30*time.Second); err != nil {
|
||||
t.Fatalf("PutInFlight: %v", err)
|
||||
}
|
||||
if err := AddPending(ctx, rdb, "orders", "0", 1, 7); err != nil {
|
||||
t.Fatalf("AddPending: %v", err)
|
||||
}
|
||||
if _, err := AcquireFIFOLock(ctx, rdb, "orders", "g1", "rh-1", 30*time.Second); err != nil {
|
||||
t.Fatalf("AcquireFIFOLock: %v", err)
|
||||
}
|
||||
|
||||
outcome, err := Ack(ctx, rdb, "orders", "rh-1")
|
||||
if err != nil || outcome != AckOutcomeAcked {
|
||||
t.Fatalf("Ack = %v, err=%v, want acked", outcome, err)
|
||||
}
|
||||
|
||||
if _, ok, _ := GetInFlight(ctx, rdb, "orders", "rh-1"); ok {
|
||||
t.Error("in-flight record should be gone after ack")
|
||||
}
|
||||
if _, ok, _ := MinPending(ctx, rdb, "orders", "0", 1); ok {
|
||||
t.Error("pending offset should be removed after ack")
|
||||
}
|
||||
locked, err := AcquireFIFOLock(ctx, rdb, "orders", "g1", "rh-2", 30*time.Second)
|
||||
if err != nil || !locked {
|
||||
t.Errorf("fifo lock should be released after ack: locked=%v err=%v", locked, err)
|
||||
}
|
||||
|
||||
// Acking again is a no-op, not an error.
|
||||
outcome, err = Ack(ctx, rdb, "orders", "rh-1")
|
||||
if err != nil || outcome != AckOutcomeNotFound {
|
||||
t.Fatalf("second Ack = %v, err=%v, want not_found", outcome, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReapRedeliversBelowMaxReceiveCount(t *testing.T) {
|
||||
rdb := newTestClient(t)
|
||||
ctx := context.Background()
|
||||
|
||||
rec := InFlightRecord{ShardID: "0", Topic: "kmsvc.orders.shard-0", Partition: 1, Offset: 7, GroupID: "g1", ReceiveCount: 1, Body: "payload"}
|
||||
if err := PutInFlight(ctx, rdb, "orders", "rh-1", rec, time.Millisecond); err != nil {
|
||||
t.Fatalf("PutInFlight: %v", err)
|
||||
}
|
||||
if _, err := AcquireFIFOLock(ctx, rdb, "orders", "g1", "rh-1", 30*time.Second); err != nil {
|
||||
t.Fatalf("AcquireFIFOLock: %v", err)
|
||||
}
|
||||
|
||||
res, err := Reap(ctx, rdb, "orders", "rh-1", 5)
|
||||
if err != nil {
|
||||
t.Fatalf("Reap: %v", err)
|
||||
}
|
||||
if res.Outcome != ReapOutcomeRedeliver || res.ReceiveCount != 2 {
|
||||
t.Fatalf("Reap = %+v, want redeliver with receiveCount=2", res)
|
||||
}
|
||||
|
||||
handle, ok, err := PopRedeliverable(ctx, rdb, "orders")
|
||||
if err != nil || !ok || handle != "rh-1" {
|
||||
t.Fatalf("PopRedeliverable = %q, ok=%v err=%v, want rh-1", handle, ok, err)
|
||||
}
|
||||
|
||||
got, ok, err := GetInFlight(ctx, rdb, "orders", "rh-1")
|
||||
if err != nil || !ok || got.ReceiveCount != 2 || got.Body != "payload" {
|
||||
t.Fatalf("GetInFlight after reap = %+v, ok=%v err=%v", got, ok, err)
|
||||
}
|
||||
|
||||
relocked, err := AcquireFIFOLock(ctx, rdb, "orders", "g1", "rh-2", 30*time.Second)
|
||||
if err != nil || !relocked {
|
||||
t.Errorf("fifo lock should be released after reap-redeliver: ok=%v err=%v", relocked, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReapRoutesToDLQAtMaxReceiveCount(t *testing.T) {
|
||||
rdb := newTestClient(t)
|
||||
ctx := context.Background()
|
||||
|
||||
rec := InFlightRecord{ShardID: "0", Topic: "kmsvc.orders.shard-0", Partition: 1, Offset: 7, GroupID: "g1", ReceiveCount: 5, Body: "payload"}
|
||||
if err := PutInFlight(ctx, rdb, "orders", "rh-1", rec, time.Millisecond); err != nil {
|
||||
t.Fatalf("PutInFlight: %v", err)
|
||||
}
|
||||
if err := AddPending(ctx, rdb, "orders", "0", 1, 7); err != nil {
|
||||
t.Fatalf("AddPending: %v", err)
|
||||
}
|
||||
|
||||
res, err := Reap(ctx, rdb, "orders", "rh-1", 5)
|
||||
if err != nil {
|
||||
t.Fatalf("Reap: %v", err)
|
||||
}
|
||||
if res.Outcome != ReapOutcomeDLQ || res.Topic != rec.Topic || res.Offset != rec.Offset || res.Body != "payload" {
|
||||
t.Fatalf("Reap = %+v, want dlq with topic/offset/body matching the record", res)
|
||||
}
|
||||
|
||||
if _, ok, _ := GetInFlight(ctx, rdb, "orders", "rh-1"); ok {
|
||||
t.Error("in-flight record should be gone after DLQ routing")
|
||||
}
|
||||
if _, ok, _ := MinPending(ctx, rdb, "orders", "0", 1); ok {
|
||||
t.Error("pending offset should be removed after DLQ routing")
|
||||
}
|
||||
}
|
||||
|
||||
func TestReapConcurrentRaceHasExactlyOneWinner(t *testing.T) {
|
||||
rdb := newTestClient(t)
|
||||
ctx := context.Background()
|
||||
|
||||
rec := InFlightRecord{ShardID: "0", Topic: "kmsvc.orders.shard-0", Partition: 1, Offset: 7, ReceiveCount: 4}
|
||||
if err := PutInFlight(ctx, rdb, "orders", "rh-1", rec, time.Millisecond); err != nil {
|
||||
t.Fatalf("PutInFlight: %v", err)
|
||||
}
|
||||
|
||||
const n = 20
|
||||
var wins atomic.Int32
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < n; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
res, err := Reap(ctx, rdb, "orders", "rh-1", 5)
|
||||
if err != nil {
|
||||
t.Errorf("Reap: %v", err)
|
||||
return
|
||||
}
|
||||
if res.Outcome != ReapOutcomeGone {
|
||||
wins.Add(1)
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
if got := wins.Load(); got != 1 {
|
||||
t.Errorf("expected exactly 1 non-gone outcome among %d concurrent reapers, got %d", n, got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,55 @@
|
||||
package redis
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
|
||||
"github.com/redis/go-redis/v9"
|
||||
|
||||
"github.com/rockliang/kafka-management-service/internal/kafka"
|
||||
)
|
||||
|
||||
// PutShardMap writes the active shard set for a queue (design.md §4's
|
||||
// `kmsvc:shardmap:` row) and publishes an invalidation so cached readers
|
||||
// refresh. Written by the queue-operator on every reconcile that changes
|
||||
// shard membership (initial create, split, close).
|
||||
func PutShardMap(ctx context.Context, rdb *redis.Client, queue string, shards []kafka.Shard) error {
|
||||
data, err := json.Marshal(shards)
|
||||
if err != nil {
|
||||
return fmt.Errorf("marshal shard map %s: %w", queue, err)
|
||||
}
|
||||
pipe := rdb.TxPipeline()
|
||||
pipe.Set(ctx, ShardMapKey(queue), data, 0)
|
||||
pipe.Publish(ctx, ShardMapChannel(queue), "invalidate")
|
||||
if _, err := pipe.Exec(ctx); err != nil {
|
||||
return fmt.Errorf("put shard map %s: %w", queue, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetShardMap reads the active shard set for a queue, or ok=false if the
|
||||
// queue has no shard map yet (not yet reconciled).
|
||||
func GetShardMap(ctx context.Context, rdb *redis.Client, queue string) ([]kafka.Shard, bool, error) {
|
||||
data, err := rdb.Get(ctx, ShardMapKey(queue)).Bytes()
|
||||
if err == redis.Nil {
|
||||
return nil, false, nil
|
||||
}
|
||||
if err != nil {
|
||||
return nil, false, fmt.Errorf("get shard map %s: %w", queue, err)
|
||||
}
|
||||
var shards []kafka.Shard
|
||||
if err := json.Unmarshal(data, &shards); err != nil {
|
||||
return nil, false, fmt.Errorf("unmarshal shard map %s: %w", queue, err)
|
||||
}
|
||||
return shards, true, nil
|
||||
}
|
||||
|
||||
// DeleteShardMap removes a queue's shard map, used by the queue-operator on
|
||||
// Queue deletion.
|
||||
func DeleteShardMap(ctx context.Context, rdb *redis.Client, queue string) error {
|
||||
if err := rdb.Del(ctx, ShardMapKey(queue)).Err(); err != nil {
|
||||
return fmt.Errorf("delete shard map %s: %w", queue, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,72 @@
|
||||
package redis
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
// AddPending records a freshly-consumed offset as not-yet-acked, per design.md
|
||||
// §3/§4's `kmsvc:pending:` ZSET.
|
||||
func AddPending(ctx context.Context, rdb *redis.Client, queue, shardID string, partition int32, offset int64) error {
|
||||
key := PendingKey(queue, shardID, partition)
|
||||
if err := rdb.ZAdd(ctx, key, redis.Z{Score: float64(offset), Member: offset}).Err(); err != nil {
|
||||
return fmt.Errorf("add pending %s: %w", key, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// MinPending returns the lowest not-yet-acked offset for a (shard, partition),
|
||||
// or ok=false if nothing is pending — i.e. the watermark can advance to the
|
||||
// last consumed offset.
|
||||
func MinPending(ctx context.Context, rdb *redis.Client, queue, shardID string, partition int32) (offset int64, ok bool, err error) {
|
||||
res, err := rdb.ZRangeWithScores(ctx, PendingKey(queue, shardID, partition), 0, 0).Result()
|
||||
if err != nil {
|
||||
return 0, false, fmt.Errorf("min pending %s: %w", PendingKey(queue, shardID, partition), err)
|
||||
}
|
||||
if len(res) == 0 {
|
||||
return 0, false, nil
|
||||
}
|
||||
return int64(res[0].Score), true, nil
|
||||
}
|
||||
|
||||
// SetWatermark stores the committable watermark for a (shard, partition),
|
||||
// per design.md §3's background committer.
|
||||
func SetWatermark(ctx context.Context, rdb *redis.Client, queue, shardID string, partition int32, offset int64) error {
|
||||
key := WatermarkKey(queue, shardID, partition)
|
||||
if err := rdb.Set(ctx, key, offset, 0).Err(); err != nil {
|
||||
return fmt.Errorf("set watermark %s: %w", key, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetWatermark returns the last-committed watermark for a (shard, partition),
|
||||
// or ok=false if none has been committed yet.
|
||||
func GetWatermark(ctx context.Context, rdb *redis.Client, queue, shardID string, partition int32) (offset int64, ok bool, err error) {
|
||||
v, err := rdb.Get(ctx, WatermarkKey(queue, shardID, partition)).Result()
|
||||
if err == redis.Nil {
|
||||
return 0, false, nil
|
||||
}
|
||||
if err != nil {
|
||||
return 0, false, fmt.Errorf("get watermark %s: %w", WatermarkKey(queue, shardID, partition), err)
|
||||
}
|
||||
n, err := strconv.ParseInt(v, 10, 64)
|
||||
if err != nil {
|
||||
return 0, false, fmt.Errorf("parse watermark %s: %w", WatermarkKey(queue, shardID, partition), err)
|
||||
}
|
||||
return n, true, nil
|
||||
}
|
||||
|
||||
// AcquireFIFOLock claims the per-group exclusivity gate (design.md §3/§4's
|
||||
// `kmsvc:fifo_lock:` row) so only one in-flight message per MessageGroupId
|
||||
// is ever handed out.
|
||||
func AcquireFIFOLock(ctx context.Context, rdb *redis.Client, queue, groupID, receiptHandle string, ttl time.Duration) (bool, error) {
|
||||
ok, err := rdb.SetNX(ctx, FIFOLockKey(queue, groupID), receiptHandle, ttl).Result()
|
||||
if err != nil {
|
||||
return false, fmt.Errorf("acquire fifo lock %s: %w", groupID, err)
|
||||
}
|
||||
return ok, nil
|
||||
}
|
||||
Reference in New Issue
Block a user