From 3ac4083a68aabc79e78d1f62756523a303ba27f1 Mon Sep 17 00:00:00 2001 From: riotpiaole <19826264+Riotpiaole@users.noreply.github.com> Date: Sun, 21 Jun 2026 17:08:17 -0700 Subject: [PATCH] feat(operator): add queue-operator CRD reconciler with shard split/drain Reconciles Queue CRs into Kafka topics + Redis shard-map/queue-meta state: creates shard-0 on first reconcile, splits a shard's hash range into two children once its split threshold is crossed, drains and closes a parent shard once its consumer group has fully caught up and its retention window has elapsed, and tears down every shard's topic + Redis state on deletion. Includes the manager entrypoint (cmd/queue-operator) and RBAC. --- cmd/queue-operator/main.go | 111 +++++++++ config/rbac/role.yaml | 20 ++ config/rbac/role_binding.yaml | 12 + internal/operator/admin.go | 16 ++ internal/operator/fake_admin.go | 71 ++++++ internal/operator/queue_controller.go | 212 ++++++++++++++++ internal/operator/queue_controller_test.go | 275 +++++++++++++++++++++ internal/operator/shard_drain.go | 48 ++++ internal/operator/shard_split.go | 86 +++++++ 9 files changed, 851 insertions(+) create mode 100644 cmd/queue-operator/main.go create mode 100644 config/rbac/role.yaml create mode 100644 config/rbac/role_binding.yaml create mode 100644 internal/operator/admin.go create mode 100644 internal/operator/fake_admin.go create mode 100644 internal/operator/queue_controller.go create mode 100644 internal/operator/queue_controller_test.go create mode 100644 internal/operator/shard_drain.go create mode 100644 internal/operator/shard_split.go diff --git a/cmd/queue-operator/main.go b/cmd/queue-operator/main.go new file mode 100644 index 0000000..36910e3 --- /dev/null +++ b/cmd/queue-operator/main.go @@ -0,0 +1,111 @@ +// Command queue-operator runs the Queue CRD reconciler described in +// design.md §2a/§2c (shard-aware control plane: creates shard topics, +// performs capacity-driven shard splits, drains closed shards, and +// publishes queue metadata + the shard map to Redis for the message plane). +package main + +import ( + "context" + "fmt" + "os" + "strconv" + "strings" + "time" + + goredis "github.com/redis/go-redis/v9" + "k8s.io/apimachinery/pkg/runtime" + clientgoscheme "k8s.io/client-go/kubernetes/scheme" + ctrl "sigs.k8s.io/controller-runtime" + "sigs.k8s.io/controller-runtime/pkg/log/zap" + "sigs.k8s.io/controller-runtime/pkg/reconcile" + + kmsvcv1 "github.com/rockliang/kafka-management-service/apis/kmsvc/v1" + "github.com/rockliang/kafka-management-service/internal/kafka" + "github.com/rockliang/kafka-management-service/internal/operator" +) + +func main() { + ctrl.SetLogger(zap.New()) + + brokers := splitCSV(getEnv("KMSVC_KAFKA_BROKERS", "localhost:9092")) + redisAddr := getEnv("KMSVC_REDIS_ADDR", "localhost:6379") + redisPassword := os.Getenv("KMSVC_REDIS_PASSWORD") + redisDB := 0 + if v := os.Getenv("KMSVC_REDIS_DB"); v != "" { + n, err := strconv.Atoi(v) + if err != nil { + exitf("invalid KMSVC_REDIS_DB: %v", err) + } + redisDB = n + } + + scheme := runtime.NewScheme() + if err := clientgoscheme.AddToScheme(scheme); err != nil { + exitf("registering client-go scheme: %v", err) + } + if err := kmsvcv1.AddToScheme(scheme); err != nil { + exitf("registering kmsvc scheme: %v", err) + } + + mgr, err := ctrl.NewManager(ctrl.GetConfigOrDie(), ctrl.Options{Scheme: scheme}) + if err != nil { + exitf("creating manager: %v", err) + } + + admin, err := kafka.NewAdmin(brokers) + if err != nil { + exitf("creating kafka admin: %v", err) + } + + rdb := goredis.NewClient(&goredis.Options{Addr: redisAddr, Password: redisPassword, DB: redisDB}) + if err := rdb.Ping(context.Background()).Err(); err != nil { + exitf("connecting to redis at %s: %v", redisAddr, err) + } + + reconciler := &operator.QueueReconciler{ + Client: mgr.GetClient(), + Admin: admin, + Redis: rdb, + Now: time.Now, + } + + err = ctrl.NewControllerManagedBy(mgr). + For(&kmsvcv1.Queue{}). + Complete(reconcile.Func(func(ctx context.Context, req reconcile.Request) (reconcile.Result, error) { + if err := reconciler.Reconcile(ctx, req.Name); err != nil { + return reconcile.Result{}, err + } + return reconcile.Result{}, nil + })) + if err != nil { + exitf("setting up controller: %v", err) + } + + if err := mgr.Start(ctrl.SetupSignalHandler()); err != nil { + exitf("running manager: %v", err) + } +} + +func getEnv(key, def string) string { + if v := os.Getenv(key); v != "" { + return v + } + return def +} + +func splitCSV(s string) []string { + parts := strings.Split(s, ",") + out := make([]string, 0, len(parts)) + for _, p := range parts { + p = strings.TrimSpace(p) + if p != "" { + out = append(out, p) + } + } + return out +} + +func exitf(format string, args ...any) { + fmt.Fprintf(os.Stderr, format+"\n", args...) + os.Exit(1) +} diff --git a/config/rbac/role.yaml b/config/rbac/role.yaml new file mode 100644 index 0000000..d37137f --- /dev/null +++ b/config/rbac/role.yaml @@ -0,0 +1,20 @@ +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRole +metadata: + name: queue-operator +rules: + - apiGroups: ["kmsvc.io"] + resources: ["queues"] + verbs: ["get", "list", "watch", "update", "patch"] + - apiGroups: ["kmsvc.io"] + resources: ["queues/status"] + verbs: ["get", "update", "patch"] + - apiGroups: ["kmsvc.io"] + resources: ["queues/finalizers"] + verbs: ["update"] + - apiGroups: ["coordination.k8s.io"] + resources: ["leases"] + verbs: ["get", "list", "watch", "create", "update", "patch", "delete"] + - apiGroups: [""] + resources: ["events"] + verbs: ["create", "patch"] diff --git a/config/rbac/role_binding.yaml b/config/rbac/role_binding.yaml new file mode 100644 index 0000000..e78609d --- /dev/null +++ b/config/rbac/role_binding.yaml @@ -0,0 +1,12 @@ +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRoleBinding +metadata: + name: queue-operator +roleRef: + apiGroup: rbac.authorization.k8s.io + kind: ClusterRole + name: queue-operator +subjects: + - kind: ServiceAccount + name: queue-operator + namespace: kmsvc diff --git a/internal/operator/admin.go b/internal/operator/admin.go new file mode 100644 index 0000000..0111222 --- /dev/null +++ b/internal/operator/admin.go @@ -0,0 +1,16 @@ +package operator + +import ( + "context" + + "github.com/rockliang/kafka-management-service/internal/kafka" +) + +// TopicAdmin is the subset of internal/kafka.Admin the reconciler needs, +// abstracted so tests can substitute a fake instead of a real Kafka cluster. +type TopicAdmin interface { + CreateTopic(ctx context.Context, topic string, cfg kafka.TopicConfig) error + DeleteTopic(ctx context.Context, topic string) error + LogEndOffsetSum(ctx context.Context, topic string) (int64, error) + ConsumerLag(ctx context.Context, group, topic string) (int64, error) +} diff --git a/internal/operator/fake_admin.go b/internal/operator/fake_admin.go new file mode 100644 index 0000000..1dc8f2f --- /dev/null +++ b/internal/operator/fake_admin.go @@ -0,0 +1,71 @@ +package operator + +import ( + "context" + "sync" + + "github.com/rockliang/kafka-management-service/internal/kafka" +) + +// fakeAdmin is an in-memory TopicAdmin for reconciler tests — avoids needing +// a real Kafka cluster (or envtest's apiserver binaries, which aren't +// available in this environment) just to exercise reconcile logic. +type fakeAdmin struct { + mu sync.Mutex + topics map[string]kafka.TopicConfig + offsetSums map[string]int64 + lag map[string]int64 // keyed by group+"/"+topic +} + +func newFakeAdmin() *fakeAdmin { + return &fakeAdmin{ + topics: map[string]kafka.TopicConfig{}, + offsetSums: map[string]int64{}, + lag: map[string]int64{}, + } +} + +func (f *fakeAdmin) CreateTopic(_ context.Context, topic string, cfg kafka.TopicConfig) error { + f.mu.Lock() + defer f.mu.Unlock() + f.topics[topic] = cfg + return nil +} + +func (f *fakeAdmin) DeleteTopic(_ context.Context, topic string) error { + f.mu.Lock() + defer f.mu.Unlock() + delete(f.topics, topic) + return nil +} + +func (f *fakeAdmin) LogEndOffsetSum(_ context.Context, topic string) (int64, error) { + f.mu.Lock() + defer f.mu.Unlock() + return f.offsetSums[topic], nil +} + +func (f *fakeAdmin) ConsumerLag(_ context.Context, group, topic string) (int64, error) { + f.mu.Lock() + defer f.mu.Unlock() + return f.lag[group+"/"+topic], nil +} + +func (f *fakeAdmin) setOffsetSum(topic string, v int64) { + f.mu.Lock() + defer f.mu.Unlock() + f.offsetSums[topic] = v +} + +func (f *fakeAdmin) setLag(group, topic string, v int64) { + f.mu.Lock() + defer f.mu.Unlock() + f.lag[group+"/"+topic] = v +} + +func (f *fakeAdmin) hasTopic(topic string) bool { + f.mu.Lock() + defer f.mu.Unlock() + _, ok := f.topics[topic] + return ok +} diff --git a/internal/operator/queue_controller.go b/internal/operator/queue_controller.go new file mode 100644 index 0000000..5bf1ad4 --- /dev/null +++ b/internal/operator/queue_controller.go @@ -0,0 +1,212 @@ +// Package operator implements the queue-operator control-plane reconciler +// described in design.md §2a/§2c: it turns a Queue CRD into one-or-more Kafka +// shard topics plus the Redis queue-metadata/shard-map the message-plane +// service reads on its hot path. +package operator + +import ( + "context" + "fmt" + "strconv" + "time" + + goredis "github.com/redis/go-redis/v9" + apierrors "k8s.io/apimachinery/pkg/api/errors" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "sigs.k8s.io/controller-runtime/pkg/client" + "sigs.k8s.io/controller-runtime/pkg/controller/controllerutil" + + kmsvcv1 "github.com/rockliang/kafka-management-service/apis/kmsvc/v1" + "github.com/rockliang/kafka-management-service/internal/kafka" + kmsvcredis "github.com/rockliang/kafka-management-service/internal/redis" +) + +const finalizerName = "kmsvc.io/queue-operator" + +// avgMessageSizeBytesEstimate converts the record-rate sampled from Kafka +// end-offset deltas into an approximate bytes/sec figure to compare against +// ShardSplitThresholdBytesPerSec, since there's no metrics pipeline (e.g. +// Prometheus) wired up in v1 to get a real byte rate per design.md §2c. +const avgMessageSizeBytesEstimate = 1024 + +const replicationFactor = 3 +const minInsyncReplicas = 2 + +// QueueReconciler reconciles Queue objects (design.md §2a). +type QueueReconciler struct { + Client client.Client + Admin TopicAdmin + Redis *goredis.Client + Now func() time.Time + + // sampleState tracks the last (offsetSum, time) seen per shard topic, used + // to compute a throughput estimate between reconciles. Keyed by topic name. + sampleState map[string]sample +} + +type sample struct { + offsetSum int64 + at time.Time +} + +func (r *QueueReconciler) now() time.Time { + if r.Now != nil { + return r.Now() + } + return time.Now() +} + +// Reconcile implements the controller-runtime reconcile loop. +func (r *QueueReconciler) Reconcile(ctx context.Context, name string) error { + var queue kmsvcv1.Queue + err := r.Client.Get(ctx, client.ObjectKey{Name: name}, &queue) + if apierrors.IsNotFound(err) { + return nil + } + if err != nil { + return fmt.Errorf("get queue %s: %w", name, err) + } + + if !queue.DeletionTimestamp.IsZero() { + return r.reconcileDelete(ctx, &queue) + } + + if err := kafka.ValidateNoDLQCycle(queue.Name, queue.Spec.IsDLQ, queue.Spec.DeadLetterTargetQueue); err != nil { + return r.setFailed(ctx, &queue, "DLQCycle", err) + } + + if !controllerutil.ContainsFinalizer(&queue, finalizerName) { + controllerutil.AddFinalizer(&queue, finalizerName) + if err := r.Client.Update(ctx, &queue); err != nil { + return fmt.Errorf("add finalizer %s: %w", name, err) + } + } + + if len(queue.Status.Shards) == 0 { + queue.Status.Shards = []kmsvcv1.ShardStatus{r.newShard("0", "", 0, kafka.FullHashRangeEnd, &queue)} + } + + if err := r.ensureShardTopics(ctx, &queue); err != nil { + return r.setFailed(ctx, &queue, "EnsureTopics", err) + } + + if err := r.reconcileSplits(ctx, &queue); err != nil { + return r.setFailed(ctx, &queue, "ShardSplit", err) + } + + if err := r.reconcileDrains(ctx, &queue); err != nil { + return r.setFailed(ctx, &queue, "ShardDrain", err) + } + + if err := r.publishRedisState(ctx, &queue); err != nil { + return r.setFailed(ctx, &queue, "PublishRedis", err) + } + + queue.Status.Phase = kmsvcv1.QueuePhaseReady + if err := r.Client.Status().Update(ctx, &queue); err != nil { + return fmt.Errorf("update status %s: %w", name, err) + } + return nil +} + +func (r *QueueReconciler) newShard(id, parentID string, start, end uint32, queue *kmsvcv1.Queue) kmsvcv1.ShardStatus { + return kmsvcv1.ShardStatus{ + ID: id, + Topic: kafka.ShardTopicName(queue.Name, queue.Spec.FIFOQueue, id), + HashRangeStart: start, + HashRangeEnd: end, + Phase: kmsvcv1.ShardPhaseActive, + ParentID: parentID, + CreatedAt: metav1.NewTime(r.now()), + } +} + +func (r *QueueReconciler) ensureShardTopics(ctx context.Context, queue *kmsvcv1.Queue) error { + cfg := kafka.TopicConfig{ + PartitionCount: queue.Spec.PartitionsPerShard, + ReplicationFactor: replicationFactor, + RetentionSeconds: queue.Spec.MessageRetentionPeriodSeconds, + MinInsyncReplicas: minInsyncReplicas, + } + for _, s := range queue.Status.Shards { + if s.Phase == kmsvcv1.ShardPhaseClosed { + continue + } + if err := r.Admin.CreateTopic(ctx, s.Topic, cfg); err != nil { + return fmt.Errorf("ensure topic %s: %w", s.Topic, err) + } + } + return nil +} + +func (r *QueueReconciler) setFailed(ctx context.Context, queue *kmsvcv1.Queue, reason string, cause error) error { + queue.Status.Phase = kmsvcv1.QueuePhaseFailed + if err := r.Client.Status().Update(ctx, queue); err != nil { + return fmt.Errorf("update failed status %s (reason=%s, cause=%v): %w", queue.Name, reason, cause, err) + } + return fmt.Errorf("reconcile %s failed (%s): %w", queue.Name, reason, cause) +} + +func (r *QueueReconciler) publishRedisState(ctx context.Context, queue *kmsvcv1.Queue) error { + if err := kmsvcredis.PutQueueMeta(ctx, r.Redis, queue.Name, kmsvcredis.QueueMeta{ + FIFO: queue.Spec.FIFOQueue, + VisibilityTimeoutSeconds: queue.Spec.VisibilityTimeoutSeconds, + MaxReceiveCount: queue.Spec.MaxReceiveCount, + DLQQueueName: queue.Spec.DeadLetterTargetQueue, + PartitionsPerShard: queue.Spec.PartitionsPerShard, + RetentionSeconds: queue.Spec.MessageRetentionPeriodSeconds, + CreatedAt: queue.CreationTimestamp.Time, + }); err != nil { + return err + } + + shards := make([]kafka.Shard, 0, len(queue.Status.Shards)) + for _, s := range queue.Status.Shards { + if s.Phase == kmsvcv1.ShardPhaseClosed { + continue + } + shards = append(shards, kafka.Shard{ + ID: s.ID, + Topic: s.Topic, + HashRangeStart: s.HashRangeStart, + HashRangeEnd: s.HashRangeEnd, + Phase: string(s.Phase), + }) + } + return kmsvcredis.PutShardMap(ctx, r.Redis, queue.Name, shards) +} + +func (r *QueueReconciler) reconcileDelete(ctx context.Context, queue *kmsvcv1.Queue) error { + if !controllerutil.ContainsFinalizer(queue, finalizerName) { + return nil + } + for _, s := range queue.Status.Shards { + if s.Phase == kmsvcv1.ShardPhaseClosed { + continue + } + if err := r.Admin.DeleteTopic(ctx, s.Topic); err != nil { + return fmt.Errorf("delete topic %s: %w", s.Topic, err) + } + } + if err := kmsvcredis.DeleteQueueMeta(ctx, r.Redis, queue.Name); err != nil { + return err + } + if err := kmsvcredis.DeleteShardMap(ctx, r.Redis, queue.Name); err != nil { + return err + } + controllerutil.RemoveFinalizer(queue, finalizerName) + if err := r.Client.Update(ctx, queue); err != nil { + return fmt.Errorf("remove finalizer %s: %w", queue.Name, err) + } + return nil +} + +func nextShardID(shards []kmsvcv1.ShardStatus) string { + max := -1 + for _, s := range shards { + if n, err := strconv.Atoi(s.ID); err == nil && n > max { + max = n + } + } + return strconv.Itoa(max + 1) +} diff --git a/internal/operator/queue_controller_test.go b/internal/operator/queue_controller_test.go new file mode 100644 index 0000000..b96b149 --- /dev/null +++ b/internal/operator/queue_controller_test.go @@ -0,0 +1,275 @@ +package operator + +import ( + "context" + "testing" + "time" + + "github.com/alicebob/miniredis/v2" + goredis "github.com/redis/go-redis/v9" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/runtime" + clientgoscheme "k8s.io/client-go/kubernetes/scheme" + "sigs.k8s.io/controller-runtime/pkg/client" + "sigs.k8s.io/controller-runtime/pkg/client/fake" + + kmsvcv1 "github.com/rockliang/kafka-management-service/apis/kmsvc/v1" + "github.com/rockliang/kafka-management-service/internal/kafka" + kmsvcredis "github.com/rockliang/kafka-management-service/internal/redis" +) + +func newTestScheme(t *testing.T) *runtime.Scheme { + t.Helper() + scheme := runtime.NewScheme() + if err := clientgoscheme.AddToScheme(scheme); err != nil { + t.Fatalf("add client-go scheme: %v", err) + } + if err := kmsvcv1.AddToScheme(scheme); err != nil { + t.Fatalf("add kmsvc scheme: %v", err) + } + return scheme +} + +func newTestRedis(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 newTestReconciler(t *testing.T, objs ...client.Object) (*QueueReconciler, *fakeAdmin) { + t.Helper() + scheme := newTestScheme(t) + cl := fake.NewClientBuilder(). + WithScheme(scheme). + WithStatusSubresource(&kmsvcv1.Queue{}). + WithObjects(objs...). + Build() + admin := newFakeAdmin() + return &QueueReconciler{ + Client: cl, + Admin: admin, + Redis: newTestRedis(t), + Now: time.Now, + }, admin +} + +func baseQueue(name string) *kmsvcv1.Queue { + return &kmsvcv1.Queue{ + ObjectMeta: metav1.ObjectMeta{Name: name}, + Spec: kmsvcv1.QueueSpec{ + VisibilityTimeoutSeconds: 30, + MessageRetentionPeriodSeconds: 345600, + MaxReceiveCount: 5, + PartitionsPerShard: 6, + MinShards: 1, + MaxShards: 8, + }, + } +} + +func TestReconcileCreatesShardZeroAndPublishesRedisState(t *testing.T) { + queue := baseQueue("orders") + r, admin := newTestReconciler(t, queue) + ctx := context.Background() + + if err := r.Reconcile(ctx, "orders"); err != nil { + t.Fatalf("Reconcile: %v", err) + } + + var got kmsvcv1.Queue + if err := r.Client.Get(ctx, client.ObjectKey{Name: "orders"}, &got); err != nil { + t.Fatalf("get queue: %v", err) + } + if got.Status.Phase != kmsvcv1.QueuePhaseReady { + t.Fatalf("phase = %v, want Ready", got.Status.Phase) + } + if len(got.Status.Shards) != 1 || got.Status.Shards[0].ID != "0" { + t.Fatalf("shards = %+v, want one shard-0", got.Status.Shards) + } + wantTopic := kafka.ShardTopicName("orders", false, "0") + if got.Status.Shards[0].Topic != wantTopic { + t.Errorf("shard-0 topic = %q, want %q", got.Status.Shards[0].Topic, wantTopic) + } + if !admin.hasTopic(wantTopic) { + t.Errorf("expected topic %q to be created", wantTopic) + } + + meta, ok, err := kmsvcredis.GetQueueMeta(ctx, r.Redis, "orders") + if err != nil || !ok { + t.Fatalf("GetQueueMeta: ok=%v err=%v", ok, err) + } + if meta.MaxReceiveCount != 5 { + t.Errorf("meta.MaxReceiveCount = %d, want 5", meta.MaxReceiveCount) + } + + shards, ok, err := kmsvcredis.GetShardMap(ctx, r.Redis, "orders") + if err != nil || !ok || len(shards) != 1 { + t.Fatalf("GetShardMap: shards=%+v ok=%v err=%v", shards, ok, err) + } +} + +func TestReconcileRejectsDLQSelfReference(t *testing.T) { + queue := baseQueue("orders") + queue.Spec.DeadLetterTargetQueue = "orders" + r, _ := newTestReconciler(t, queue) + ctx := context.Background() + + err := r.Reconcile(ctx, "orders") + if err == nil { + t.Fatal("expected reconcile to fail for self-referencing DLQ") + } + + var got kmsvcv1.Queue + if err := r.Client.Get(ctx, client.ObjectKey{Name: "orders"}, &got); err != nil { + t.Fatalf("get queue: %v", err) + } + if got.Status.Phase != kmsvcv1.QueuePhaseFailed { + t.Errorf("phase = %v, want Failed", got.Status.Phase) + } +} + +func TestReconcileDeleteCleansUpTopicsAndRedis(t *testing.T) { + queue := baseQueue("orders") + r, admin := newTestReconciler(t, queue) + ctx := context.Background() + + if err := r.Reconcile(ctx, "orders"); err != nil { + t.Fatalf("initial reconcile: %v", err) + } + + var got kmsvcv1.Queue + if err := r.Client.Get(ctx, client.ObjectKey{Name: "orders"}, &got); err != nil { + t.Fatalf("get queue: %v", err) + } + topic := got.Status.Shards[0].Topic + + if err := r.Client.Delete(ctx, &got); err != nil { + t.Fatalf("delete queue: %v", err) + } + if err := r.Reconcile(ctx, "orders"); err != nil { + t.Fatalf("delete reconcile: %v", err) + } + + if admin.hasTopic(topic) { + t.Errorf("expected topic %q to be deleted", topic) + } + if _, ok, _ := kmsvcredis.GetQueueMeta(ctx, r.Redis, "orders"); ok { + t.Error("expected queue meta to be removed") + } + if _, ok, _ := kmsvcredis.GetShardMap(ctx, r.Redis, "orders"); ok { + t.Error("expected shard map to be removed") + } +} + +func TestReconcileSplitsShardOverThreshold(t *testing.T) { + queue := baseQueue("orders") + queue.Spec.ShardSplitThresholdBytesPerSec = 1000 + queue.Spec.ShardSplitCooldownSeconds = 0 + now := time.Now() + r, admin := newTestReconciler(t, queue) + r.Now = func() time.Time { return now } + ctx := context.Background() + + if err := r.Reconcile(ctx, "orders"); err != nil { + t.Fatalf("initial reconcile: %v", err) + } + var got kmsvcv1.Queue + if err := r.Client.Get(ctx, client.ObjectKey{Name: "orders"}, &got); err != nil { + t.Fatalf("get queue: %v", err) + } + shard0Topic := got.Status.Shards[0].Topic + admin.setOffsetSum(shard0Topic, 0) + + // First sample establishes the baseline (no delta to compare yet). + if err := r.Reconcile(ctx, "orders"); err != nil { + t.Fatalf("baseline reconcile: %v", err) + } + + // Advance time and offsets enough to exceed 1000 bytes/sec at the + // 1024-byte/record estimate: 10 records over 1s ≈ 10240 bytes/sec. + now = now.Add(time.Second) + admin.setOffsetSum(shard0Topic, 10) + + if err := r.Reconcile(ctx, "orders"); err != nil { + t.Fatalf("split-triggering reconcile: %v", err) + } + + if err := r.Client.Get(ctx, client.ObjectKey{Name: "orders"}, &got); err != nil { + t.Fatalf("get queue after split: %v", err) + } + if len(got.Status.Shards) != 3 { + t.Fatalf("shards = %+v, want 3 (1 closing parent + 2 active children)", got.Status.Shards) + } + var parent *kmsvcv1.ShardStatus + var children []kmsvcv1.ShardStatus + for i := range got.Status.Shards { + s := &got.Status.Shards[i] + if s.ID == "0" { + parent = s + } else { + children = append(children, *s) + } + } + if parent == nil || parent.Phase != kmsvcv1.ShardPhaseClosing { + t.Fatalf("parent shard = %+v, want Closing", parent) + } + if len(children) != 2 { + t.Fatalf("children = %+v, want 2", children) + } + mid := kafka.SplitHashRange(parent.HashRangeStart, parent.HashRangeEnd) + gotRanges := map[[2]uint32]bool{} + for _, c := range children { + if c.Phase != kmsvcv1.ShardPhaseActive { + t.Errorf("child %s phase = %v, want Active", c.ID, c.Phase) + } + gotRanges[[2]uint32{c.HashRangeStart, c.HashRangeEnd}] = true + } + if !gotRanges[[2]uint32{0, mid}] || !gotRanges[[2]uint32{mid, kafka.FullHashRangeEnd}] { + t.Errorf("child ranges = %+v, want [0,%d) and [%d,%d)", children, mid, mid, kafka.FullHashRangeEnd) + } +} + +func TestReconcileDrainsClosingShardWhenLagZero(t *testing.T) { + queue := baseQueue("orders") + queue.Spec.MessageRetentionPeriodSeconds = 60 + now := time.Now() + r, admin := newTestReconciler(t, queue) + r.Now = func() time.Time { return now } + ctx := context.Background() + + if err := r.Reconcile(ctx, "orders"); err != nil { + t.Fatalf("initial reconcile: %v", err) + } + var got kmsvcv1.Queue + if err := r.Client.Get(ctx, client.ObjectKey{Name: "orders"}, &got); err != nil { + t.Fatalf("get queue: %v", err) + } + + // Manually mark shard-0 as Closing (simulating a prior split) to isolate + // drain behavior from split behavior. + got.Status.Shards[0].Phase = kmsvcv1.ShardPhaseClosing + got.Status.Shards[0].CreatedAt = metav1.NewTime(now.Add(-2 * time.Minute)) + if err := r.Client.Status().Update(ctx, &got); err != nil { + t.Fatalf("seed closing shard: %v", err) + } + topic := got.Status.Shards[0].Topic + admin.setLag(kafka.ConsumerGroup("orders"), topic, 0) + + if err := r.Reconcile(ctx, "orders"); err != nil { + t.Fatalf("drain reconcile: %v", err) + } + + if err := r.Client.Get(ctx, client.ObjectKey{Name: "orders"}, &got); err != nil { + t.Fatalf("get queue after drain: %v", err) + } + if got.Status.Shards[0].Phase != kmsvcv1.ShardPhaseClosed { + t.Errorf("shard-0 phase = %v, want Closed", got.Status.Shards[0].Phase) + } + if admin.hasTopic(topic) { + t.Errorf("expected drained topic %q to be deleted", topic) + } +} diff --git a/internal/operator/shard_drain.go b/internal/operator/shard_drain.go new file mode 100644 index 0000000..4d3b7fa --- /dev/null +++ b/internal/operator/shard_drain.go @@ -0,0 +1,48 @@ +package operator + +import ( + "context" + "fmt" + "time" + + kmsvcv1 "github.com/rockliang/kafka-management-service/apis/kmsvc/v1" + "github.com/rockliang/kafka-management-service/internal/kafka" +) + +func secondsToDuration(s int32) time.Duration { + return time.Duration(s) * time.Second +} + +// reconcileDrains transitions `Closing` shards to `Closed` once their +// consumer-group lag reaches zero and the retention window has elapsed since +// they stopped being a write target, then deletes their topic, per +// design.md §2c. +func (r *QueueReconciler) reconcileDrains(ctx context.Context, queue *kmsvcv1.Queue) error { + group := kafka.ConsumerGroup(queue.Name) + retention := secondsToDuration(queue.Spec.MessageRetentionPeriodSeconds) + + kept := make([]kmsvcv1.ShardStatus, 0, len(queue.Status.Shards)) + for _, s := range queue.Status.Shards { + if s.Phase != kmsvcv1.ShardPhaseClosing { + kept = append(kept, s) + continue + } + + lag, err := r.Admin.ConsumerLag(ctx, group, s.Topic) + if err != nil { + return fmt.Errorf("consumer lag %s: %w", s.Topic, err) + } + if lag > 0 || r.now().Before(s.CreatedAt.Add(retention)) { + kept = append(kept, s) + continue + } + + if err := r.Admin.DeleteTopic(ctx, s.Topic); err != nil { + return fmt.Errorf("delete drained topic %s: %w", s.Topic, err) + } + s.Phase = kmsvcv1.ShardPhaseClosed + kept = append(kept, s) + } + queue.Status.Shards = kept + return nil +} diff --git a/internal/operator/shard_split.go b/internal/operator/shard_split.go new file mode 100644 index 0000000..9a84aca --- /dev/null +++ b/internal/operator/shard_split.go @@ -0,0 +1,86 @@ +package operator + +import ( + "context" + "fmt" + "strconv" + + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + + kmsvcv1 "github.com/rockliang/kafka-management-service/apis/kmsvc/v1" + "github.com/rockliang/kafka-management-service/internal/kafka" +) + +// reconcileSplits samples each Active shard's throughput and splits any shard +// that's sustained over spec.ShardSplitThresholdBytesPerSec into two children, +// per design.md §2c. +func (r *QueueReconciler) reconcileSplits(ctx context.Context, queue *kmsvcv1.Queue) error { + if r.sampleState == nil { + r.sampleState = map[string]sample{} + } + + maxShards := queue.Spec.MaxShards + if maxShards <= 0 { + maxShards = 8 + } + threshold := queue.Spec.ShardSplitThresholdBytesPerSec + if threshold <= 0 { + return nil + } + cooldown := queue.Spec.ShardSplitCooldownSeconds + + activeCount := 0 + for _, s := range queue.Status.Shards { + if s.Phase == kmsvcv1.ShardPhaseActive { + activeCount++ + } + } + + for i := range queue.Status.Shards { + s := &queue.Status.Shards[i] + if s.Phase != kmsvcv1.ShardPhaseActive { + continue + } + if int32(activeCount) >= maxShards { + break + } + if !s.CreatedAt.IsZero() && r.now().Before(s.CreatedAt.Add(secondsToDuration(cooldown))) { + continue + } + + offsetSum, err := r.Admin.LogEndOffsetSum(ctx, s.Topic) + if err != nil { + return fmt.Errorf("sample throughput %s: %w", s.Topic, err) + } + prev, seen := r.sampleState[s.Topic] + now := r.now() + r.sampleState[s.Topic] = sample{offsetSum: offsetSum, at: now} + if !seen { + continue + } + elapsed := now.Sub(prev.at).Seconds() + if elapsed <= 0 { + continue + } + recordsPerSec := float64(offsetSum-prev.offsetSum) / elapsed + bytesPerSec := recordsPerSec * avgMessageSizeBytesEstimate + if bytesPerSec < float64(threshold) { + continue + } + + mid := kafka.SplitHashRange(s.HashRangeStart, s.HashRangeEnd) + childAID := nextShardID(queue.Status.Shards) + childA := r.newShard(childAID, s.ID, s.HashRangeStart, mid, queue) + childA.CreatedAt = metav1.NewTime(now) + childAIDNum, _ := strconv.Atoi(childAID) + childBID := strconv.Itoa(childAIDNum + 1) + childB := r.newShard(childBID, s.ID, mid, s.HashRangeEnd, queue) + childB.CreatedAt = metav1.NewTime(now) + + s.Phase = kmsvcv1.ShardPhaseClosing + queue.Status.Shards = append(queue.Status.Shards, childA, childB) + activeCount += 1 // net: -1 parent +2 children + delete(r.sampleState, s.Topic) + } + return nil +}