kmsvc-manage/internal/kafka/admin.go

package kafka

import (
	"context"
	"fmt"
	"strconv"

	"github.com/twmb/franz-go/pkg/kadm"
	"github.com/twmb/franz-go/pkg/kgo"
)

// Admin wraps kadm.Client with the topic conventions from design.md §6.
type Admin struct {
	client *kadm.Client
}

// NewAdmin builds an Admin from a set of broker addresses.
func NewAdmin(brokers []string) (*Admin, error) {
	cl, err := kgo.NewClient(kgo.SeedBrokers(brokers...))
	if err != nil {
		return nil, fmt.Errorf("creating kafka client: %w", err)
	}
	return &Admin{client: kadm.NewClient(cl)}, nil
}

// TopicConfig is the subset of Kafka topic configuration design.md §6 cares about.
type TopicConfig struct {
	PartitionCount    int32
	ReplicationFactor int16
	RetentionSeconds  int32
	MinInsyncReplicas int16
}

// CreateTopic creates a topic idempotently: if it already exists with the
// requested partition count, this is a no-op (design.md §3 task 3
// acceptance criteria).
func (a *Admin) CreateTopic(ctx context.Context, topic string, cfg TopicConfig) error {
	configs := map[string]*string{
		"retention.ms":        strPtr(strconv.FormatInt(int64(cfg.RetentionSeconds)*1000, 10)),
		"cleanup.policy":      strPtr("delete"),
		"min.insync.replicas": strPtr(strconv.Itoa(int(cfg.MinInsyncReplicas))),
	}

	resp, err := a.client.CreateTopics(ctx, cfg.PartitionCount, cfg.ReplicationFactor, configs, topic)
	if err != nil {
		return fmt.Errorf("create topic %q: %w", topic, err)
	}
	for _, t := range resp {
		if t.Err != nil && !isTopicExistsErr(t.Err) {
			return fmt.Errorf("create topic %q: %w", topic, t.Err)
		}
	}
	return nil
}

// DeleteTopic deletes a topic. Deleting a non-existent topic is not an error.
func (a *Admin) DeleteTopic(ctx context.Context, topic string) error {
	resp, err := a.client.DeleteTopics(ctx, topic)
	if err != nil {
		return fmt.Errorf("delete topic %q: %w", topic, err)
	}
	for _, t := range resp {
		if t.Err != nil && !isUnknownTopicErr(t.Err) {
			return fmt.Errorf("delete topic %q: %w", topic, t.Err)
		}
	}
	return nil
}

// DescribeTopic returns the live partition count and config for a topic, or
// ok=false if it does not exist.
func (a *Admin) DescribeTopic(ctx context.Context, topic string) (partitions int32, ok bool, err error) {
	td, err := a.client.ListTopics(ctx, topic)
	if err != nil {
		return 0, false, fmt.Errorf("describe topic %q: %w", topic, err)
	}
	detail, found := td[topic]
	if !found || detail.Err != nil {
		return 0, false, nil
	}
	return int32(len(detail.Partitions)), true, nil
}

// ReplicaBrokerIDs returns the set of broker IDs holding any replica of any
// partition of topic, used by the queue-operator to resolve which
// availability zones (node labels) a shard's data actually lives on
// (design.md §2a AZ-awareness).
func (a *Admin) ReplicaBrokerIDs(ctx context.Context, topic string) ([]int32, error) {
	td, err := a.client.ListTopics(ctx, topic)
	if err != nil {
		return nil, fmt.Errorf("list topics %q: %w", topic, err)
	}
	detail, found := td[topic]
	if !found || detail.Err != nil {
		return nil, nil
	}

	seen := make(map[int32]bool)
	var ids []int32
	for _, p := range detail.Partitions {
		for _, id := range p.Replicas {
			if !seen[id] {
				seen[id] = true
				ids = append(ids, id)
			}
		}
	}
	return ids, nil
}

// LogEndOffsetSum returns the sum of the high-watermark offsets across every
// partition of a topic — a monotonically increasing proxy for total records
// produced, used by the queue-operator's shard-split sampler (design.md §2c)
// to estimate throughput between two reconcile ticks.
func (a *Admin) LogEndOffsetSum(ctx context.Context, topic string) (int64, error) {
	offsets, err := a.client.ListEndOffsets(ctx, topic)
	if err != nil {
		return 0, fmt.Errorf("list end offsets %q: %w", topic, err)
	}
	var sum int64
	for _, partitions := range offsets {
		for _, o := range partitions {
			if o.Err != nil {
				continue
			}
			sum += o.Offset
		}
	}
	return sum, nil
}

// ConsumerLag returns the total lag (sum across partitions) of a consumer
// group against a topic, used by the queue-operator to decide when a
// `Closing` shard (design.md §2c) is fully drained.
func (a *Admin) ConsumerLag(ctx context.Context, group, topic string) (int64, error) {
	lags, err := a.client.Lag(ctx, group)
	if err != nil {
		return 0, fmt.Errorf("lag for group %q: %w", group, err)
	}
	described, ok := lags[group]
	if !ok || described.Error() != nil {
		return 0, nil
	}
	return described.Lag.TotalByTopic()[topic].Lag, nil
}

// CommitOffset commits the next offset to read for (group, topic, partition)
// to Kafka's real consumer-group offsets — the message plane calls this on
// ack to advance the low-watermark commit described in design.md §3, rather
// than relying on the consumer client's own interval autocommit (which is
// disabled, see internal/kafka.Consumer).
func (a *Admin) CommitOffset(ctx context.Context, group, topic string, partition int32, offset int64) error {
	offsets := kadm.Offsets{topic: {partition: kadm.Offset{Topic: topic, Partition: partition, At: offset}}}
	resp, err := a.client.CommitOffsets(ctx, group, offsets)
	if err != nil {
		return fmt.Errorf("commit offset %s/%s/%d: %w", group, topic, partition, err)
	}
	if err := resp.Error(); err != nil {
		return fmt.Errorf("commit offset %s/%s/%d: %w", group, topic, partition, err)
	}
	return nil
}

func (a *Admin) Close() {
	a.client.Close()
}

func strPtr(s string) *string { return &s }

func isTopicExistsErr(err error) bool {
	return err != nil && (err.Error() == "TOPIC_ALREADY_EXISTS" || containsCode(err, "TOPIC_ALREADY_EXISTS"))
}

func isUnknownTopicErr(err error) bool {
	return err != nil && containsCode(err, "UNKNOWN_TOPIC_OR_PARTITION")
}

func containsCode(err error, code string) bool {
	return err != nil && (err.Error() == code || len(err.Error()) >= len(code) && (err.Error()[:len(code)] == code))
}