diff --git a/internal/redis/client.go b/internal/redis/client.go new file mode 100644 index 0000000..32cda88 --- /dev/null +++ b/internal/redis/client.go @@ -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, + }) +} diff --git a/internal/redis/dedup.go b/internal/redis/dedup.go new file mode 100644 index 0000000..02294f0 --- /dev/null +++ b/internal/redis/dedup.go @@ -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 +} diff --git a/internal/redis/inflight.go b/internal/redis/inflight.go new file mode 100644 index 0000000..5e05293 --- /dev/null +++ b/internal/redis/inflight.go @@ -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 +} diff --git a/internal/redis/keys.go b/internal/redis/keys.go new file mode 100644 index 0000000..32994f8 --- /dev/null +++ b/internal/redis/keys.go @@ -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) +} diff --git a/internal/redis/keys_test.go b/internal/redis/keys_test.go new file mode 100644 index 0000000..aa03476 --- /dev/null +++ b/internal/redis/keys_test.go @@ -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) + } + }) + } +} diff --git a/internal/redis/lua/ack.lua b/internal/redis/lua/ack.lua new file mode 100644 index 0000000..39c0363 --- /dev/null +++ b/internal/redis/lua/ack.lua @@ -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'} diff --git a/internal/redis/lua/reap.lua b/internal/redis/lua/reap.lua new file mode 100644 index 0000000..7a24045 --- /dev/null +++ b/internal/redis/lua/reap.lua @@ -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)} diff --git a/internal/redis/queuemeta.go b/internal/redis/queuemeta.go new file mode 100644 index 0000000..8c26c3e --- /dev/null +++ b/internal/redis/queuemeta.go @@ -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 +} diff --git a/internal/redis/reaper.go b/internal/redis/reaper.go new file mode 100644 index 0000000..081cf18 --- /dev/null +++ b/internal/redis/reaper.go @@ -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 +} diff --git a/internal/redis/redis_test.go b/internal/redis/redis_test.go new file mode 100644 index 0000000..b0ef9ff --- /dev/null +++ b/internal/redis/redis_test.go @@ -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) + } +} diff --git a/internal/redis/shardmap.go b/internal/redis/shardmap.go new file mode 100644 index 0000000..5903e51 --- /dev/null +++ b/internal/redis/shardmap.go @@ -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 +} diff --git a/internal/redis/watermark.go b/internal/redis/watermark.go new file mode 100644 index 0000000..3e950b4 --- /dev/null +++ b/internal/redis/watermark.go @@ -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 +}