From b1f039ae2554acb6c6487c174f721cd56f093c17 Mon Sep 17 00:00:00 2001 From: riotpiaole <19826264+Riotpiaole@users.noreply.github.com> Date: Sun, 21 Jun 2026 17:08:24 -0700 Subject: [PATCH] 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. --- internal/core/queue/change_visibility.go | 29 +++ internal/core/queue/delete_message.go | 75 ++++++ internal/core/queue/fifo.go | 25 ++ internal/core/queue/queue_test.go | 306 +++++++++++++++++++++++ internal/core/queue/receive_message.go | 226 +++++++++++++++++ internal/core/queue/send_message.go | 110 ++++++++ internal/core/queue/shard_router.go | 71 ++++++ 7 files changed, 842 insertions(+) create mode 100644 internal/core/queue/change_visibility.go create mode 100644 internal/core/queue/delete_message.go create mode 100644 internal/core/queue/fifo.go create mode 100644 internal/core/queue/queue_test.go create mode 100644 internal/core/queue/receive_message.go create mode 100644 internal/core/queue/send_message.go create mode 100644 internal/core/queue/shard_router.go diff --git a/internal/core/queue/change_visibility.go b/internal/core/queue/change_visibility.go new file mode 100644 index 0000000..8ec318f --- /dev/null +++ b/internal/core/queue/change_visibility.go @@ -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 +} diff --git a/internal/core/queue/delete_message.go b/internal/core/queue/delete_message.go new file mode 100644 index 0000000..f8efe39 --- /dev/null +++ b/internal/core/queue/delete_message.go @@ -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 +} diff --git a/internal/core/queue/fifo.go b/internal/core/queue/fifo.go new file mode 100644 index 0000000..2d06a2d --- /dev/null +++ b/internal/core/queue/fifo.go @@ -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 +} diff --git a/internal/core/queue/queue_test.go b/internal/core/queue/queue_test.go new file mode 100644 index 0000000..1aa72c8 --- /dev/null +++ b/internal/core/queue/queue_test.go @@ -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") + } +} diff --git a/internal/core/queue/receive_message.go b/internal/core/queue/receive_message.go new file mode 100644 index 0000000..2189690 --- /dev/null +++ b/internal/core/queue/receive_message.go @@ -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()) +} diff --git a/internal/core/queue/send_message.go b/internal/core/queue/send_message.go new file mode 100644 index 0000000..4398b52 --- /dev/null +++ b/internal/core/queue/send_message.go @@ -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 +} diff --git a/internal/core/queue/shard_router.go b/internal/core/queue/shard_router.go new file mode 100644 index 0000000..64c913f --- /dev/null +++ b/internal/core/queue/shard_router.go @@ -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 +}