feat(queue): add message-plane core logic for send/receive/delete/visibility
Implements SendMessage (size cap, FIFO dedup, shard routing), ReceiveMessage (long-poll loop across active+closing shards, FIFO per-group exclusivity gating, redelivery hand-out), DeleteMessage (ack + low-watermark offset advancement), and ChangeMessageVisibility against the Redis state layer and a real Kafka producer/consumer.
This commit is contained in:
riotpiaole
@@ -0,0 +1,29 @@
|
||||
package queue
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
goredis "github.com/redis/go-redis/v9"
|
||||
|
||||
kmsvcredis "github.com/rockliang/kafka-management-service/internal/redis"
|
||||
)
|
||||
|
||||
type ChangeVisibilityService struct {
|
||||
Redis *goredis.Client
|
||||
}
|
||||
|
||||
// ChangeMessageVisibility implements ChangeMessageVisibility (design.md
|
||||
// §2b): extends (or shortens) how long a received message stays invisible
|
||||
// to other consumers, without affecting its receive count.
|
||||
func (s *ChangeVisibilityService) ChangeMessageVisibility(ctx context.Context, queueName, receiptHandle string, newTimeout time.Duration) error {
|
||||
ok, err := kmsvcredis.ExtendVisibility(ctx, s.Redis, queueName, receiptHandle, newTimeout)
|
||||
if err != nil {
|
||||
return fmt.Errorf("change visibility %s: %w", queueName, err)
|
||||
}
|
||||
if !ok {
|
||||
return fmt.Errorf("receipt handle not found or already expired")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,75 @@
|
||||
package queue
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
goredis "github.com/redis/go-redis/v9"
|
||||
|
||||
"github.com/rockliang/kafka-management-service/internal/kafka"
|
||||
kmsvcredis "github.com/rockliang/kafka-management-service/internal/redis"
|
||||
)
|
||||
|
||||
// OffsetCommitter is the subset of kafka.Admin DeleteMessage needs to
|
||||
// advance a (shard, partition)'s committed consumer-group offset, per
|
||||
// design.md §3's low-watermark strategy.
|
||||
type OffsetCommitter interface {
|
||||
CommitOffset(ctx context.Context, group, topic string, partition int32, offset int64) error
|
||||
}
|
||||
|
||||
type DeleteMessageService struct {
|
||||
Redis *goredis.Client
|
||||
Committer OffsetCommitter
|
||||
}
|
||||
|
||||
// DeleteMessage acks a received message (design.md §3): the atomic ack.lua
|
||||
// script removes its in-flight record + vis_index entry, removes its offset
|
||||
// from the pending set, and releases any FIFO lock it held; this then
|
||||
// advances the (shard, partition)'s committed watermark if doing so is safe.
|
||||
// Calling DeleteMessage on an already-acked/DLQ-routed/redelivered receipt
|
||||
// handle is a no-op, matching SQS's idempotent DeleteMessage semantics.
|
||||
func (s *DeleteMessageService) DeleteMessage(ctx context.Context, queueName, receiptHandle string) error {
|
||||
rec, ok, err := kmsvcredis.GetInFlight(ctx, s.Redis, queueName, receiptHandle)
|
||||
if err != nil {
|
||||
return fmt.Errorf("delete message %s: %w", queueName, err)
|
||||
}
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
|
||||
outcome, err := kmsvcredis.Ack(ctx, s.Redis, queueName, receiptHandle)
|
||||
if err != nil {
|
||||
return fmt.Errorf("delete message %s: %w", queueName, err)
|
||||
}
|
||||
if outcome != kmsvcredis.AckOutcomeAcked {
|
||||
return nil
|
||||
}
|
||||
|
||||
return s.advanceWatermark(ctx, queueName, rec)
|
||||
}
|
||||
|
||||
// advanceWatermark commits the new low-watermark for the acked message's
|
||||
// (shard, partition) once the pending set has a new minimum. If nothing is
|
||||
// pending (everything ever consumed on this partition has been acked), the
|
||||
// last commit is left as-is rather than invented from nothing — at worst
|
||||
// this means a restart replays slightly further back than strictly
|
||||
// necessary, which design.md §3 accepts as the at-least-once tradeoff.
|
||||
func (s *DeleteMessageService) advanceWatermark(ctx context.Context, queueName string, rec kmsvcredis.InFlightRecord) error {
|
||||
minOffset, ok, err := kmsvcredis.MinPending(ctx, s.Redis, queueName, rec.ShardID, rec.Partition)
|
||||
if err != nil {
|
||||
return fmt.Errorf("advance watermark %s: %w", queueName, err)
|
||||
}
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
if err := kmsvcredis.SetWatermark(ctx, s.Redis, queueName, rec.ShardID, rec.Partition, minOffset-1); err != nil {
|
||||
return fmt.Errorf("advance watermark %s: %w", queueName, err)
|
||||
}
|
||||
if s.Committer == nil {
|
||||
return nil
|
||||
}
|
||||
if err := s.Committer.CommitOffset(ctx, kafka.ConsumerGroup(queueName), rec.Topic, rec.Partition, minOffset); err != nil {
|
||||
return fmt.Errorf("advance watermark %s: %w", queueName, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,25 @@
|
||||
package queue
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
goredis "github.com/redis/go-redis/v9"
|
||||
|
||||
kmsvcredis "github.com/rockliang/kafka-management-service/internal/redis"
|
||||
)
|
||||
|
||||
// acquireFIFOSlot claims the per-group exclusivity gate (design.md §3) so
|
||||
// ReceiveMessage never hands out two in-flight messages for the same
|
||||
// MessageGroupId at once. The lock's TTL matches the message's own
|
||||
// visibility timeout, so a lock left by a crashed/timed-out consumer
|
||||
// self-heals on the same schedule the reaper uses (design.md §5) — it is
|
||||
// also released early, by ack.lua/reap.lua, on ack/redeliver/DLQ-route.
|
||||
func acquireFIFOSlot(ctx context.Context, rdb *goredis.Client, queueName, groupID, receiptHandle string, visibilityTimeout time.Duration) (bool, error) {
|
||||
ok, err := kmsvcredis.AcquireFIFOLock(ctx, rdb, queueName, groupID, receiptHandle, visibilityTimeout)
|
||||
if err != nil {
|
||||
return false, fmt.Errorf("acquire fifo slot %s/%s: %w", queueName, groupID, err)
|
||||
}
|
||||
return ok, nil
|
||||
}
|
||||
@@ -0,0 +1,306 @@
|
||||
package queue
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/alicebob/miniredis/v2"
|
||||
goredis "github.com/redis/go-redis/v9"
|
||||
"github.com/twmb/franz-go/pkg/kfake"
|
||||
|
||||
"github.com/rockliang/kafka-management-service/internal/kafka"
|
||||
kmsvcredis "github.com/rockliang/kafka-management-service/internal/redis"
|
||||
)
|
||||
|
||||
// newTestKafka starts an in-memory, wire-protocol-compatible fake Kafka
|
||||
// cluster (kfake) so these tests exercise the real franz-go producer/
|
||||
// consumer/admin code paths without needing Docker/testcontainers — not
|
||||
// available in this environment, same documented tradeoff as the
|
||||
// queue-operator's envtest substitution (internal/operator).
|
||||
func newTestKafka(t *testing.T) []string {
|
||||
t.Helper()
|
||||
cluster, err := kfake.NewCluster(kfake.NumBrokers(1))
|
||||
if err != nil {
|
||||
t.Fatalf("starting kfake cluster: %v", err)
|
||||
}
|
||||
t.Cleanup(cluster.Close)
|
||||
return cluster.ListenAddrs()
|
||||
}
|
||||
|
||||
func newTestRedisClient(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()})
|
||||
}
|
||||
|
||||
// seedQueue creates the queue's shard-0 topic, queue meta, and shard map
|
||||
// directly — standing in for the queue-operator (internal/operator), which
|
||||
// is exercised separately in its own test suite.
|
||||
func seedQueue(t *testing.T, ctx context.Context, brokers []string, rdb *goredis.Client, name string, fifo bool, partitionsPerShard int32) kafka.Shard {
|
||||
t.Helper()
|
||||
admin, err := kafka.NewAdmin(brokers)
|
||||
if err != nil {
|
||||
t.Fatalf("new admin: %v", err)
|
||||
}
|
||||
t.Cleanup(admin.Close)
|
||||
|
||||
shard := kafka.Shard{ID: "0", Topic: kafka.ShardTopicName(name, fifo, "0"), HashRangeStart: 0, HashRangeEnd: kafka.FullHashRangeEnd, Phase: "Active"}
|
||||
if err := admin.CreateTopic(ctx, shard.Topic, kafka.TopicConfig{PartitionCount: partitionsPerShard, ReplicationFactor: 1, RetentionSeconds: 345600, MinInsyncReplicas: 1}); err != nil {
|
||||
t.Fatalf("create shard topic: %v", err)
|
||||
}
|
||||
if err := kmsvcredis.PutQueueMeta(ctx, rdb, name, kmsvcredis.QueueMeta{
|
||||
FIFO: fifo,
|
||||
VisibilityTimeoutSeconds: 1,
|
||||
MaxReceiveCount: 3,
|
||||
PartitionsPerShard: partitionsPerShard,
|
||||
RetentionSeconds: 345600,
|
||||
}); err != nil {
|
||||
t.Fatalf("put queue meta: %v", err)
|
||||
}
|
||||
if err := kmsvcredis.PutShardMap(ctx, rdb, name, []kafka.Shard{shard}); err != nil {
|
||||
t.Fatalf("put shard map: %v", err)
|
||||
}
|
||||
return shard
|
||||
}
|
||||
|
||||
func newTestServices(t *testing.T, brokers []string, rdb *goredis.Client, group string, topics ...string) (*SendMessageService, *ReceiveMessageService, *DeleteMessageService) {
|
||||
t.Helper()
|
||||
producer, err := kafka.NewProducer(brokers)
|
||||
if err != nil {
|
||||
t.Fatalf("new producer: %v", err)
|
||||
}
|
||||
t.Cleanup(producer.Close)
|
||||
|
||||
consumer, err := kafka.NewConsumer(brokers, group, topics...)
|
||||
if err != nil {
|
||||
t.Fatalf("new consumer: %v", err)
|
||||
}
|
||||
t.Cleanup(consumer.Close)
|
||||
|
||||
router := &ShardRouter{Redis: rdb}
|
||||
send := &SendMessageService{Redis: rdb, Producer: producer, Router: router}
|
||||
receive := &ReceiveMessageService{Redis: rdb, Fetcher: consumer, Router: router, PollInterval: 20 * time.Millisecond}
|
||||
del := &DeleteMessageService{Redis: rdb}
|
||||
return send, receive, del
|
||||
}
|
||||
|
||||
func receiveUntil(t *testing.T, ctx context.Context, receive *ReceiveMessageService, queueName string, want int, waitTime time.Duration) []Message {
|
||||
t.Helper()
|
||||
out, err := receive.ReceiveMessage(ctx, ReceiveMessageInput{QueueName: queueName, MaxNumberOfMessages: int32(want), WaitTime: waitTime})
|
||||
if err != nil {
|
||||
t.Fatalf("receive message: %v", err)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func TestSendReceiveDeleteStandardQueue(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
brokers := newTestKafka(t)
|
||||
rdb := newTestRedisClient(t)
|
||||
shard := seedQueue(t, ctx, brokers, rdb, "orders", false, 1)
|
||||
send, receive, del := newTestServices(t, brokers, rdb, kafka.ConsumerGroup("orders"), shard.Topic)
|
||||
|
||||
if _, err := send.SendMessage(ctx, SendMessageInput{QueueName: "orders", Body: "hello"}); err != nil {
|
||||
t.Fatalf("send message: %v", err)
|
||||
}
|
||||
|
||||
msgs := receiveUntil(t, ctx, receive, "orders", 1, 2*time.Second)
|
||||
if len(msgs) != 1 || msgs[0].Body != "hello" {
|
||||
t.Fatalf("messages = %+v, want one %q", msgs, "hello")
|
||||
}
|
||||
|
||||
if err := del.DeleteMessage(ctx, "orders", msgs[0].ReceiptHandle); err != nil {
|
||||
t.Fatalf("delete message: %v", err)
|
||||
}
|
||||
|
||||
// Nothing left to redeliver: a short poll should come back empty.
|
||||
again := receiveUntil(t, ctx, receive, "orders", 1, 100*time.Millisecond)
|
||||
if len(again) != 0 {
|
||||
t.Fatalf("messages after delete = %+v, want none", again)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReceiveMessageRedeliversAfterVisibilityTimeoutExpires(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
brokers := newTestKafka(t)
|
||||
rdb := newTestRedisClient(t)
|
||||
shard := seedQueue(t, ctx, brokers, rdb, "orders", false, 1)
|
||||
send, receive, _ := newTestServices(t, brokers, rdb, kafka.ConsumerGroup("orders"), shard.Topic)
|
||||
|
||||
if _, err := send.SendMessage(ctx, SendMessageInput{QueueName: "orders", Body: "redeliver-me"}); err != nil {
|
||||
t.Fatalf("send message: %v", err)
|
||||
}
|
||||
|
||||
first := receiveUntil(t, ctx, receive, "orders", 1, 2*time.Second)
|
||||
if len(first) != 1 {
|
||||
t.Fatalf("first receive = %+v, want one message", first)
|
||||
}
|
||||
|
||||
// Run the reaper directly (its own service is covered in internal/redis
|
||||
// and a separate Task 7 reaper loop — here we just need its effect: an
|
||||
// expired in-flight message becomes redeliverable).
|
||||
if _, err := kmsvcredis.Reap(ctx, rdb, "orders", first[0].ReceiptHandle, 3); err != nil {
|
||||
t.Fatalf("reap: %v", err)
|
||||
}
|
||||
|
||||
second := receiveUntil(t, ctx, receive, "orders", 1, 2*time.Second)
|
||||
if len(second) != 1 || second[0].ReceiveCount != 2 {
|
||||
t.Fatalf("second receive = %+v, want one message with receiveCount=2", second)
|
||||
}
|
||||
if second[0].Body != "redeliver-me" {
|
||||
t.Fatalf("redelivered body = %q, want %q", second[0].Body, "redeliver-me")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFIFOSecondMessageInGroupNotReceivableUntilFirstAcked(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
brokers := newTestKafka(t)
|
||||
rdb := newTestRedisClient(t)
|
||||
shard := seedQueue(t, ctx, brokers, rdb, "orders.fifo", true, 1)
|
||||
send, receive, del := newTestServices(t, brokers, rdb, kafka.ConsumerGroup("orders.fifo"), shard.Topic)
|
||||
|
||||
if _, err := send.SendMessage(ctx, SendMessageInput{QueueName: "orders.fifo", Body: "first", MessageGroupID: "g1"}); err != nil {
|
||||
t.Fatalf("send first: %v", err)
|
||||
}
|
||||
if _, err := send.SendMessage(ctx, SendMessageInput{QueueName: "orders.fifo", Body: "second", MessageGroupID: "g1"}); err != nil {
|
||||
t.Fatalf("send second: %v", err)
|
||||
}
|
||||
|
||||
first := receiveUntil(t, ctx, receive, "orders.fifo", 2, 2*time.Second)
|
||||
if len(first) != 1 || first[0].Body != "first" {
|
||||
t.Fatalf("first receive = %+v, want only %q", first, "first")
|
||||
}
|
||||
|
||||
// The second message is consumed off the partition but parked (FIFO gate
|
||||
// held) — it must not be receivable yet.
|
||||
blocked := receiveUntil(t, ctx, receive, "orders.fifo", 1, 200*time.Millisecond)
|
||||
if len(blocked) != 0 {
|
||||
t.Fatalf("messages while group locked = %+v, want none", blocked)
|
||||
}
|
||||
|
||||
if err := del.DeleteMessage(ctx, "orders.fifo", first[0].ReceiptHandle); err != nil {
|
||||
t.Fatalf("delete first: %v", err)
|
||||
}
|
||||
|
||||
second := receiveUntil(t, ctx, receive, "orders.fifo", 1, 2*time.Second)
|
||||
if len(second) != 1 || second[0].Body != "second" {
|
||||
t.Fatalf("second receive after ack = %+v, want %q", second, "second")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSendMessageDedupRejectsDuplicateWithinWindow(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
brokers := newTestKafka(t)
|
||||
rdb := newTestRedisClient(t)
|
||||
shard := seedQueue(t, ctx, brokers, rdb, "orders.fifo", true, 1)
|
||||
send, _, _ := newTestServices(t, brokers, rdb, kafka.ConsumerGroup("orders.fifo"), shard.Topic)
|
||||
|
||||
first, err := send.SendMessage(ctx, SendMessageInput{QueueName: "orders.fifo", Body: "v1", MessageGroupID: "g1", MessageDeduplicationID: "d1"})
|
||||
if err != nil {
|
||||
t.Fatalf("send first: %v", err)
|
||||
}
|
||||
dup, err := send.SendMessage(ctx, SendMessageInput{QueueName: "orders.fifo", Body: "v2 (should be ignored)", MessageGroupID: "g1", MessageDeduplicationID: "d1"})
|
||||
if err != nil {
|
||||
t.Fatalf("send dup: %v", err)
|
||||
}
|
||||
if dup.MessageID != "d1" {
|
||||
t.Errorf("dup message id = %q, want the dedup id %q", dup.MessageID, "d1")
|
||||
}
|
||||
if first.SequenceNumber == dup.SequenceNumber && first.MessageID == dup.MessageID {
|
||||
t.Errorf("expected dup response to be distinguishable from the original send")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSendMessageRejectsBodyOverSizeCap(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
brokers := newTestKafka(t)
|
||||
rdb := newTestRedisClient(t)
|
||||
shard := seedQueue(t, ctx, brokers, rdb, "orders", false, 1)
|
||||
send, _, _ := newTestServices(t, brokers, rdb, kafka.ConsumerGroup("orders"), shard.Topic)
|
||||
|
||||
oversized := make([]byte, MaxMessageBodyBytes+1)
|
||||
_, err := send.SendMessage(ctx, SendMessageInput{QueueName: "orders", Body: string(oversized)})
|
||||
if err == nil {
|
||||
t.Fatal("expected oversized message body to be rejected")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFIFOGroupStableAcrossShardSplit(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
brokers := newTestKafka(t)
|
||||
rdb := newTestRedisClient(t)
|
||||
|
||||
admin, err := kafka.NewAdmin(brokers)
|
||||
if err != nil {
|
||||
t.Fatalf("new admin: %v", err)
|
||||
}
|
||||
t.Cleanup(admin.Close)
|
||||
|
||||
queueName := "orders.fifo"
|
||||
parent := kafka.Shard{ID: "0", Topic: kafka.ShardTopicName(queueName, true, "0"), HashRangeStart: 0, HashRangeEnd: kafka.FullHashRangeEnd, Phase: "Active"}
|
||||
for _, topic := range []string{parent.Topic} {
|
||||
if err := admin.CreateTopic(ctx, topic, kafka.TopicConfig{PartitionCount: 1, ReplicationFactor: 1, RetentionSeconds: 345600, MinInsyncReplicas: 1}); err != nil {
|
||||
t.Fatalf("create topic %s: %v", topic, err)
|
||||
}
|
||||
}
|
||||
if err := kmsvcredis.PutQueueMeta(ctx, rdb, queueName, kmsvcredis.QueueMeta{FIFO: true, VisibilityTimeoutSeconds: 30, MaxReceiveCount: 3, PartitionsPerShard: 1}); err != nil {
|
||||
t.Fatalf("put queue meta: %v", err)
|
||||
}
|
||||
if err := kmsvcredis.PutShardMap(ctx, rdb, queueName, []kafka.Shard{parent}); err != nil {
|
||||
t.Fatalf("put shard map: %v", err)
|
||||
}
|
||||
|
||||
router := &ShardRouter{Redis: rdb}
|
||||
groupID := "group-A"
|
||||
beforeShard, err := router.RouteForSend(ctx, queueName, groupID)
|
||||
if err != nil {
|
||||
t.Fatalf("route before split: %v", err)
|
||||
}
|
||||
if beforeShard.ID != "0" {
|
||||
t.Fatalf("beforeShard = %+v, want shard-0", beforeShard)
|
||||
}
|
||||
|
||||
// Simulate the operator splitting shard-0 (internal/operator/shard_split.go).
|
||||
mid := kafka.SplitHashRange(parent.HashRangeStart, parent.HashRangeEnd)
|
||||
childA := kafka.Shard{ID: "1", Topic: kafka.ShardTopicName(queueName, true, "1"), HashRangeStart: 0, HashRangeEnd: mid, Phase: "Active"}
|
||||
childB := kafka.Shard{ID: "2", Topic: kafka.ShardTopicName(queueName, true, "2"), HashRangeStart: mid, HashRangeEnd: kafka.FullHashRangeEnd, Phase: "Active"}
|
||||
parent.Phase = "Closing"
|
||||
for _, topic := range []string{childA.Topic, childB.Topic} {
|
||||
if err := admin.CreateTopic(ctx, topic, kafka.TopicConfig{PartitionCount: 1, ReplicationFactor: 1, RetentionSeconds: 345600, MinInsyncReplicas: 1}); err != nil {
|
||||
t.Fatalf("create topic %s: %v", topic, err)
|
||||
}
|
||||
}
|
||||
if err := kmsvcredis.PutShardMap(ctx, rdb, queueName, []kafka.Shard{parent, childA, childB}); err != nil {
|
||||
t.Fatalf("put post-split shard map: %v", err)
|
||||
}
|
||||
|
||||
afterShard, err := router.RouteForSend(ctx, queueName, groupID)
|
||||
if err != nil {
|
||||
t.Fatalf("route after split: %v", err)
|
||||
}
|
||||
h := kafka.HashKey(groupID)
|
||||
wantID := "1"
|
||||
if h >= mid {
|
||||
wantID = "2"
|
||||
}
|
||||
if afterShard.ID != wantID {
|
||||
t.Fatalf("afterShard = %+v, want shard %s (hash=%d, mid=%d)", afterShard, wantID, h, mid)
|
||||
}
|
||||
|
||||
// ReceiveMessage consuming across both the still-draining parent and the
|
||||
// new child must still see every message, regardless of which shard it
|
||||
// landed on.
|
||||
send, receive, _ := newTestServices(t, brokers, rdb, kafka.ConsumerGroup(queueName), parent.Topic, childA.Topic, childB.Topic)
|
||||
if _, err := send.SendMessage(ctx, SendMessageInput{QueueName: queueName, Body: "post-split", MessageGroupID: groupID}); err != nil {
|
||||
t.Fatalf("send post-split: %v", err)
|
||||
}
|
||||
msgs := receiveUntil(t, ctx, receive, queueName, 1, 2*time.Second)
|
||||
if len(msgs) != 1 || msgs[0].Body != "post-split" {
|
||||
t.Fatalf("messages = %+v, want one %q", msgs, "post-split")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,226 @@
|
||||
package queue
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
goredis "github.com/redis/go-redis/v9"
|
||||
|
||||
"github.com/rockliang/kafka-management-service/internal/kafka"
|
||||
kmsvcredis "github.com/rockliang/kafka-management-service/internal/redis"
|
||||
)
|
||||
|
||||
const defaultPollInterval = 200 * time.Millisecond // design.md §2b
|
||||
|
||||
// Fetcher is the subset of kafka.Consumer ReceiveMessage needs: one
|
||||
// non-blocking poll iteration across whatever shard topics it's currently
|
||||
// subscribed to (design.md §6 — every Active+Closing shard).
|
||||
type Fetcher interface {
|
||||
Poll(ctx context.Context) ([]kafka.Record, error)
|
||||
}
|
||||
|
||||
// Message is one message handed back to a ReceiveMessage caller.
|
||||
type Message struct {
|
||||
ReceiptHandle string
|
||||
Body string
|
||||
ReceiveCount int32
|
||||
}
|
||||
|
||||
type ReceiveMessageInput struct {
|
||||
QueueName string
|
||||
MaxNumberOfMessages int32
|
||||
WaitTime time.Duration
|
||||
VisibilityTimeoutOverride time.Duration // 0 = use the queue's configured default
|
||||
}
|
||||
|
||||
type ReceiveMessageService struct {
|
||||
Redis *goredis.Client
|
||||
Fetcher Fetcher
|
||||
Router *ShardRouter
|
||||
|
||||
// PollInterval is the short sleep between Kafka+Redis poll iterations
|
||||
// during a long-poll wait (design.md §2b); defaults to 200ms.
|
||||
PollInterval time.Duration
|
||||
}
|
||||
|
||||
// ReceiveMessage implements ReceiveMessage's SQS-style long polling
|
||||
// (design.md §2b): pops anything already queued for redelivery first
|
||||
// (design.md §4, avoids re-fetching by offset), then polls Kafka for new
|
||||
// records, looping until max_number_of_messages is satisfied or wait_time
|
||||
// elapses.
|
||||
func (s *ReceiveMessageService) ReceiveMessage(ctx context.Context, in ReceiveMessageInput) ([]Message, error) {
|
||||
meta, ok, err := kmsvcredis.GetQueueMeta(ctx, s.Redis, in.QueueName)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("receive message %s: %w", in.QueueName, err)
|
||||
}
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("queue %s not found", in.QueueName)
|
||||
}
|
||||
|
||||
visTimeout := in.VisibilityTimeoutOverride
|
||||
if visTimeout <= 0 {
|
||||
visTimeout = time.Duration(meta.VisibilityTimeoutSeconds) * time.Second
|
||||
}
|
||||
max := in.MaxNumberOfMessages
|
||||
if max <= 0 || max > 10 {
|
||||
max = 10
|
||||
}
|
||||
interval := s.PollInterval
|
||||
if interval <= 0 {
|
||||
interval = defaultPollInterval
|
||||
}
|
||||
|
||||
shards, err := s.Router.ConsumableShards(ctx, in.QueueName)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
topicToShardID := make(map[string]string, len(shards))
|
||||
for _, sh := range shards {
|
||||
topicToShardID[sh.Topic] = sh.ID
|
||||
}
|
||||
|
||||
deadline := time.Now().Add(in.WaitTime)
|
||||
out := make([]Message, 0, max)
|
||||
for {
|
||||
for int32(len(out)) < max {
|
||||
msg, handed, err := s.handOutRedeliverable(ctx, meta, in.QueueName, visTimeout)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if !handed {
|
||||
break
|
||||
}
|
||||
out = append(out, msg)
|
||||
}
|
||||
|
||||
if int32(len(out)) < max {
|
||||
// kgo's PollFetches blocks until either records are available or
|
||||
// its context is done, so each iteration gets its own
|
||||
// short-lived context bounded by the poll interval — otherwise a
|
||||
// quiet topic would block the very first iteration for the
|
||||
// entire wait_time instead of retrying every interval.
|
||||
pollCtx, cancel := context.WithTimeout(ctx, interval)
|
||||
records, err := s.Fetcher.Poll(pollCtx)
|
||||
cancel()
|
||||
if err != nil && ctx.Err() != nil {
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
if err != nil && !errors.Is(err, context.DeadlineExceeded) {
|
||||
return nil, fmt.Errorf("receive message %s: %w", in.QueueName, err)
|
||||
}
|
||||
for _, rec := range records {
|
||||
if int32(len(out)) >= max {
|
||||
break
|
||||
}
|
||||
shardID, known := topicToShardID[rec.Topic]
|
||||
if !known {
|
||||
continue // topic belongs to a shard that's since closed
|
||||
}
|
||||
msg, handed, err := s.handOutFresh(ctx, meta, in.QueueName, shardID, rec, visTimeout)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if handed {
|
||||
out = append(out, msg)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if len(out) > 0 || time.Now().After(deadline) {
|
||||
return out, nil
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return out, ctx.Err()
|
||||
case <-time.After(interval):
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *ReceiveMessageService) handOutRedeliverable(ctx context.Context, meta kmsvcredis.QueueMeta, queueName string, visTimeout time.Duration) (Message, bool, error) {
|
||||
handle, ok, err := kmsvcredis.PopRedeliverable(ctx, s.Redis, queueName)
|
||||
if err != nil {
|
||||
return Message{}, false, fmt.Errorf("receive message %s: %w", queueName, err)
|
||||
}
|
||||
if !ok {
|
||||
return Message{}, false, nil
|
||||
}
|
||||
|
||||
rec, ok, err := kmsvcredis.GetInFlight(ctx, s.Redis, queueName, handle)
|
||||
if err != nil {
|
||||
return Message{}, false, fmt.Errorf("receive message %s: %w", queueName, err)
|
||||
}
|
||||
if !ok {
|
||||
// Acked/DLQ-routed between the push onto the redeliver list and this
|
||||
// pop — nothing to hand out, caller should keep draining the list.
|
||||
return Message{}, false, nil
|
||||
}
|
||||
|
||||
if rec.GroupID != "" {
|
||||
acquired, err := acquireFIFOSlot(ctx, s.Redis, queueName, rec.GroupID, handle, visTimeout)
|
||||
if err != nil {
|
||||
return Message{}, false, err
|
||||
}
|
||||
if !acquired {
|
||||
if err := kmsvcredis.PushRedeliverable(ctx, s.Redis, queueName, handle); err != nil {
|
||||
return Message{}, false, fmt.Errorf("receive message %s: %w", queueName, err)
|
||||
}
|
||||
return Message{}, false, nil
|
||||
}
|
||||
}
|
||||
|
||||
if _, err := kmsvcredis.ExtendVisibility(ctx, s.Redis, queueName, handle, visTimeout); err != nil {
|
||||
return Message{}, false, fmt.Errorf("receive message %s: %w", queueName, err)
|
||||
}
|
||||
return Message{ReceiptHandle: handle, Body: rec.Body, ReceiveCount: rec.ReceiveCount}, true, nil
|
||||
}
|
||||
|
||||
func (s *ReceiveMessageService) handOutFresh(ctx context.Context, meta kmsvcredis.QueueMeta, queueName, shardID string, rec kafka.Record, visTimeout time.Duration) (Message, bool, error) {
|
||||
if err := kmsvcredis.AddPending(ctx, s.Redis, queueName, shardID, rec.Partition, rec.Offset); err != nil {
|
||||
return Message{}, false, fmt.Errorf("receive message %s: %w", queueName, err)
|
||||
}
|
||||
|
||||
receiptHandle := newReceiptHandle(shardID, rec.Partition, rec.Offset)
|
||||
groupID := ""
|
||||
if meta.FIFO {
|
||||
groupID = string(rec.Key)
|
||||
}
|
||||
|
||||
if err := kmsvcredis.PutInFlight(ctx, s.Redis, queueName, receiptHandle, kmsvcredis.InFlightRecord{
|
||||
ShardID: shardID,
|
||||
Topic: rec.Topic,
|
||||
Partition: rec.Partition,
|
||||
Offset: rec.Offset,
|
||||
GroupID: groupID,
|
||||
Body: string(rec.Value),
|
||||
ReceiveCount: 1,
|
||||
}, visTimeout); err != nil {
|
||||
return Message{}, false, fmt.Errorf("receive message %s: %w", queueName, err)
|
||||
}
|
||||
|
||||
if groupID != "" {
|
||||
acquired, err := acquireFIFOSlot(ctx, s.Redis, queueName, groupID, receiptHandle, visTimeout)
|
||||
if err != nil {
|
||||
return Message{}, false, err
|
||||
}
|
||||
if !acquired {
|
||||
// Another message for this group is already checked out: this
|
||||
// one was already consumed off the partition (so it must be
|
||||
// tracked, not dropped) but can't be delivered yet — park it on
|
||||
// the redeliver list instead of handing it out now.
|
||||
if err := kmsvcredis.PushRedeliverable(ctx, s.Redis, queueName, receiptHandle); err != nil {
|
||||
return Message{}, false, fmt.Errorf("receive message %s: %w", queueName, err)
|
||||
}
|
||||
return Message{}, false, nil
|
||||
}
|
||||
}
|
||||
|
||||
return Message{ReceiptHandle: receiptHandle, Body: string(rec.Value), ReceiveCount: 1}, true, nil
|
||||
}
|
||||
|
||||
func newReceiptHandle(shardID string, partition int32, offset int64) string {
|
||||
return fmt.Sprintf("%s:%d:%d:%s", shardID, partition, offset, uuid.NewString())
|
||||
}
|
||||
@@ -0,0 +1,110 @@
|
||||
package queue
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
goredis "github.com/redis/go-redis/v9"
|
||||
|
||||
"github.com/rockliang/kafka-management-service/internal/kafka"
|
||||
kmsvcredis "github.com/rockliang/kafka-management-service/internal/redis"
|
||||
)
|
||||
|
||||
// MaxMessageBodyBytes is the SQS-compatible size cap enforced at
|
||||
// SendMessage, design.md §6.
|
||||
const MaxMessageBodyBytes = 256 * 1024
|
||||
|
||||
const defaultDedupWindow = 5 * time.Minute
|
||||
|
||||
// Producer is the subset of kafka.Producer SendMessage needs, abstracted so
|
||||
// tests can substitute a fake without a real Kafka broker.
|
||||
type Producer interface {
|
||||
Produce(ctx context.Context, topic string, partition int32, key, value []byte) (offset int64, err error)
|
||||
}
|
||||
|
||||
type SendMessageInput struct {
|
||||
QueueName string
|
||||
Body string
|
||||
MessageGroupID string // FIFO only
|
||||
MessageDeduplicationID string // FIFO only
|
||||
}
|
||||
|
||||
type SendMessageOutput struct {
|
||||
MessageID string
|
||||
SequenceNumber int64 // the Kafka offset it landed at
|
||||
}
|
||||
|
||||
type SendMessageService struct {
|
||||
Redis *goredis.Client
|
||||
Producer Producer
|
||||
Router *ShardRouter
|
||||
DedupWindow time.Duration
|
||||
}
|
||||
|
||||
// SendMessage implements SendMessage (design.md §2b/§2c/§6): enforces the
|
||||
// size cap, runs the FIFO dedup check, resolves the destination shard via
|
||||
// the routing key, and produces to that shard's topic at the
|
||||
// within-shard partition the key hashes to.
|
||||
func (s *SendMessageService) SendMessage(ctx context.Context, in SendMessageInput) (SendMessageOutput, error) {
|
||||
if len(in.Body) > MaxMessageBodyBytes {
|
||||
return SendMessageOutput{}, fmt.Errorf("message body exceeds %d bytes", MaxMessageBodyBytes)
|
||||
}
|
||||
|
||||
meta, ok, err := kmsvcredis.GetQueueMeta(ctx, s.Redis, in.QueueName)
|
||||
if err != nil {
|
||||
return SendMessageOutput{}, fmt.Errorf("send message %s: %w", in.QueueName, err)
|
||||
}
|
||||
if !ok {
|
||||
return SendMessageOutput{}, fmt.Errorf("queue %s not found", in.QueueName)
|
||||
}
|
||||
|
||||
if meta.FIFO {
|
||||
if in.MessageGroupID == "" {
|
||||
return SendMessageOutput{}, fmt.Errorf("messageGroupId is required for FIFO queue %s", in.QueueName)
|
||||
}
|
||||
if in.MessageDeduplicationID != "" {
|
||||
window := s.DedupWindow
|
||||
if window <= 0 {
|
||||
window = defaultDedupWindow
|
||||
}
|
||||
fresh, err := kmsvcredis.TryDedup(ctx, s.Redis, in.QueueName, in.MessageGroupID, in.MessageDeduplicationID, window)
|
||||
if err != nil {
|
||||
return SendMessageOutput{}, fmt.Errorf("send message %s: %w", in.QueueName, err)
|
||||
}
|
||||
if !fresh {
|
||||
// SQS returns the original send's message ID for a deduped
|
||||
// retry; v1 doesn't track that mapping, so the dedup ID
|
||||
// itself is returned as a stable, idempotent identifier.
|
||||
return SendMessageOutput{MessageID: in.MessageDeduplicationID}, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
routingKey := in.MessageGroupID
|
||||
if routingKey == "" {
|
||||
routingKey = uuid.NewString()
|
||||
}
|
||||
|
||||
shard, err := s.Router.RouteForSend(ctx, in.QueueName, routingKey)
|
||||
if err != nil {
|
||||
return SendMessageOutput{}, err
|
||||
}
|
||||
|
||||
partition := kafka.PartitionWithinShard(routingKey, meta.PartitionsPerShard)
|
||||
var key []byte
|
||||
if meta.FIFO {
|
||||
key = []byte(in.MessageGroupID)
|
||||
}
|
||||
|
||||
offset, err := s.Producer.Produce(ctx, shard.Topic, partition, key, []byte(in.Body))
|
||||
if err != nil {
|
||||
return SendMessageOutput{}, fmt.Errorf("send message %s: %w", in.QueueName, err)
|
||||
}
|
||||
|
||||
return SendMessageOutput{
|
||||
MessageID: uuid.NewString(),
|
||||
SequenceNumber: offset,
|
||||
}, nil
|
||||
}
|
||||
@@ -0,0 +1,71 @@
|
||||
// Package queue implements the message-plane business logic from
|
||||
// design.md §2b/§3: SendMessage/ReceiveMessage/DeleteMessage/
|
||||
// ChangeMessageVisibility, FIFO gating, and shard-aware routing.
|
||||
package queue
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
goredis "github.com/redis/go-redis/v9"
|
||||
|
||||
"github.com/rockliang/kafka-management-service/internal/kafka"
|
||||
kmsvcredis "github.com/rockliang/kafka-management-service/internal/redis"
|
||||
)
|
||||
|
||||
// ShardRouter resolves a routing key to a shard by reading the cached shard
|
||||
// map the queue-operator writes (design.md §4 `kmsvc:shardmap:`). It is the
|
||||
// only piece of core logic that needs to know shards exist — everything
|
||||
// downstream operates against a resolved (shard, topic) pair exactly as it
|
||||
// would against a single topic.
|
||||
//
|
||||
// In-process caching of the shard map (design.md §4's stated optimization)
|
||||
// is intentionally not implemented in v1: a direct Redis read per call is
|
||||
// simpler and still cheap relative to a Kafka round trip, and can be added
|
||||
// later without changing this type's interface.
|
||||
type ShardRouter struct {
|
||||
Redis *goredis.Client
|
||||
}
|
||||
|
||||
// RouteForSend returns the Active shard a new message with the given
|
||||
// routing key (MessageGroupId for FIFO, a random UUID for standard queues)
|
||||
// should be produced to.
|
||||
func (r *ShardRouter) RouteForSend(ctx context.Context, queueName, routingKey string) (kafka.Shard, error) {
|
||||
shards, err := r.shards(ctx, queueName)
|
||||
if err != nil {
|
||||
return kafka.Shard{}, err
|
||||
}
|
||||
shard, found := kafka.SelectShard(kafka.ActiveShards(shards), routingKey)
|
||||
if !found {
|
||||
return kafka.Shard{}, fmt.Errorf("no active shard covers the routing key for queue %s", queueName)
|
||||
}
|
||||
return shard, nil
|
||||
}
|
||||
|
||||
// ConsumableShards returns every shard a consumer should be subscribed to:
|
||||
// Active (write+read) and Closing (read-only, draining) — everything except
|
||||
// Closed, whose topic has already been deleted.
|
||||
func (r *ShardRouter) ConsumableShards(ctx context.Context, queueName string) ([]kafka.Shard, error) {
|
||||
shards, err := r.shards(ctx, queueName)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out := make([]kafka.Shard, 0, len(shards))
|
||||
for _, s := range shards {
|
||||
if s.Phase != "Closed" {
|
||||
out = append(out, s)
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (r *ShardRouter) shards(ctx context.Context, queueName string) ([]kafka.Shard, error) {
|
||||
shards, ok, err := kmsvcredis.GetShardMap(ctx, r.Redis, queueName)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("read shard map %s: %w", queueName, err)
|
||||
}
|
||||
if !ok || len(shards) == 0 {
|
||||
return nil, fmt.Errorf("queue %s has no shard map yet (not reconciled)", queueName)
|
||||
}
|
||||
return shards, nil
|
||||
}
|
||||
Reference in New Issue
Block a user