distributed_server.go

package server

import (
	"context"
	"errors"
	"fmt"
	"net"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/prometheus/client_golang/prometheus/promhttp"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/status"

	api "my-broker/api/v1"
	"my-broker/broker"
	"my-broker/logger"
	"my-broker/metrics"
	"my-broker/middleware"
)

// forwardedMDKey marks a request that was already proxied once by another
// broker, so proxies can't loop.
const forwardedMDKey = "x-broker-forwarded"

// produceRetryBudget bounds how long a produce keeps retrying while partition
// leadership is in flux (e.g. right after a broker failure).
const produceRetryBudget = 10 * time.Second

type DistributedServer struct {
	api.UnimplementedLogServer

	broker  *broker.DistributedBroker
	config  Config
	grpcSrv *grpc.Server
	httpSrv *http.Server

	rateLimiter *middleware.RateLimiter
	forwarder   *middleware.LeaderForwarder

	// Cached connections to other brokers for produce/consume proxying.
	proxyMu    sync.Mutex
	proxyConns map[string]*grpc.ClientConn // addr -> conn
}

// leaderProvider adapts the broker for the forwarding middleware: forwarded
// requests need the leader's client-facing gRPC address, not its Raft address.
type leaderProvider struct {
	b *broker.DistributedBroker
}

func (p leaderProvider) IsLeader() bool     { return p.b.IsLeader() }
func (p leaderProvider) LeaderAddr() string { return p.b.LeaderGRPCAddr() }

func NewDistributedServer(b *broker.DistributedBroker, cfg Config) *DistributedServer {
	return &DistributedServer{
		broker:      b,
		config:      cfg,
		rateLimiter: middleware.NewRateLimiter(1000, 1000),
		forwarder:   middleware.NewLeaderForwarder(leaderProvider{b}),
		proxyConns:  make(map[string]*grpc.ClientConn),
	}
}

func (s *DistributedServer) Start() error {
	// Join existing cluster if specified
	if s.config.JoinAddr != "" {
		logger.Info("joining cluster", "via", s.config.JoinAddr)
		if err := s.joinCluster(); err != nil {
			return fmt.Errorf("join cluster: %w", err)
		}
		logger.Info("joined cluster")
	}

	go s.startHTTPServer()

	bindAddr := s.config.BindAddr
	if bindAddr == "" {
		bindAddr = "0.0.0.0"
	}
	lis, err := net.Listen("tcp", bindAddr+":"+s.config.GRPCPort)
	if err != nil {
		return fmt.Errorf("listen: %w", err)
	}

	s.grpcSrv = grpc.NewServer(
		grpc.ChainUnaryInterceptor(
			middleware.RateLimitInterceptor(s.rateLimiter),
			middleware.LeaderForwardingInterceptor(s.forwarder),
		),
	)
	api.RegisterLogServer(s.grpcSrv, s)

	go func() {
		logger.Info("gRPC server ready", "port", s.config.GRPCPort)
		if err := s.grpcSrv.Serve(lis); err != nil {
			logger.Error("gRPC server stopped", "err", err)
		}
	}()

	return nil
}

func (s *DistributedServer) Stop() {
	if s.grpcSrv != nil {
		s.grpcSrv.GracefulStop()
	}
	if s.httpSrv != nil {
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		s.httpSrv.Shutdown(ctx)
	}
	if s.forwarder != nil {
		s.forwarder.Close()
	}
	if s.rateLimiter != nil {
		s.rateLimiter.Close()
	}

	s.proxyMu.Lock()
	for _, conn := range s.proxyConns {
		conn.Close()
	}
	s.proxyConns = make(map[string]*grpc.ClientConn)
	s.proxyMu.Unlock()
}

func (s *DistributedServer) startHTTPServer() {
	mux := http.NewServeMux()
	mux.Handle("/metrics", promhttp.Handler())
	mux.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		w.Write([]byte("OK"))
	})
	mux.HandleFunc("/ready", func(w http.ResponseWriter, r *http.Request) {
		switch {
		case s.broker.IsLeader():
			w.WriteHeader(http.StatusOK)
			w.Write([]byte("LEADER"))
		case s.broker.LeaderAddr() != "":
			w.WriteHeader(http.StatusOK)
			w.Write([]byte("FOLLOWER"))
		default:
			// No known controller leader: this node can't serve reliably.
			w.WriteHeader(http.StatusServiceUnavailable)
			w.Write([]byte("NOT_READY"))
		}
	})

	s.httpSrv = &http.Server{
		Addr:    ":" + s.config.MetricsPort,
		Handler: mux,
	}

	logger.Info("metrics server started", "port", s.config.MetricsPort)
	if err := s.httpSrv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
		logger.Error("HTTP server stopped", "err", err)
	}
}

// joinCluster contacts the seed node: it joins this node into the Raft
// cluster and registers its addresses in broker metadata so partitions can
// be assigned to it.
func (s *DistributedServer) joinCluster() error {
	addr := s.config.JoinAddr
	if !strings.Contains(addr, ":") {
		addr = addr + ":8080"
	}

	conn, err := grpc.Dial(addr, grpc.WithTransportCredentials(insecure.NewCredentials()))
	if err != nil {
		return fmt.Errorf("dial seed node: %w", err)
	}
	defer conn.Close()

	client := api.NewLogClient(conn)

	// The seed node may not be the controller leader yet (or at all); retry
	// briefly to ride out elections during cluster startup.
	var lastErr error
	for attempt := 0; attempt < 15; attempt++ {
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		_, lastErr = client.Join(ctx, &api.JoinRequest{
			NodeId:   s.config.NodeID,
			Addr:     s.broker.AdvertisedRaftAddr(),
			GrpcAddr: s.broker.AdvertisedGRPCAddr(),
			DataAddr: s.broker.AdvertisedDataAddr(),
		})
		cancel()
		if lastErr == nil {
			return nil
		}
		time.Sleep(2 * time.Second)
	}
	return lastErr
}

// resolvePartition picks the partition for a produce request: an explicit
// partition wins, a negative partition with a key hashes the key, and
// everything else lands on partition 0.
func (s *DistributedServer) resolvePartition(topic string, requested int32, key []byte) int {
	if requested >= 0 {
		return int(requested)
	}
	if len(key) > 0 {
		count := s.broker.GetTopicPartitionCount(topic)
		if count <= 0 {
			count = 1
		}
		return broker.PartitionForKey(key, count)
	}
	return 0
}

func isForwarded(ctx context.Context) bool {
	md, ok := metadata.FromIncomingContext(ctx)
	return ok && len(md.Get(forwardedMDKey)) > 0
}

// proxyConn returns a cached client connection to another broker.
func (s *DistributedServer) proxyConn(addr string) (*grpc.ClientConn, error) {
	s.proxyMu.Lock()
	defer s.proxyMu.Unlock()

	if conn, ok := s.proxyConns[addr]; ok {
		return conn, nil
	}
	conn, err := grpc.Dial(addr, grpc.WithTransportCredentials(insecure.NewCredentials()))
	if err != nil {
		return nil, err
	}
	s.proxyConns[addr] = conn
	return conn, nil
}

// partitionLeaderClient resolves the leader broker for a partition and
// returns a client for it. Returns an error when the leader is unknown or is
// this node itself (nothing to proxy to).
func (s *DistributedServer) partitionLeaderClient(topic string, partition int) (api.LogClient, error) {
	leaderAddr, err := s.broker.GetPartitionLeader(topic, partition)
	if err != nil {
		return nil, err
	}
	if leaderAddr == "" || leaderAddr == s.broker.AdvertisedGRPCAddr() {
		return nil, fmt.Errorf("no remote leader for %s-%d", topic, partition)
	}

	conn, err := s.proxyConn(leaderAddr)
	if err != nil {
		return nil, err
	}
	return api.NewLogClient(conn), nil
}

// isLeadershipErr reports whether the produce/consume failure means "wrong
// broker" (retryable elsewhere) as opposed to a real failure.
func isLeadershipErr(err error) bool {
	return errors.Is(err, broker.ErrNotLeader) || errors.Is(err, broker.ErrPartitionNotFound)
}

func (s *DistributedServer) Produce(ctx context.Context, req *api.ProduceRequest) (*api.ProduceResponse, error) {
	start := time.Now()
	topic := req.Topic
	if topic == "" {
		topic = "default"
	}

	metrics.GRPCRequestsTotal.WithLabelValues("Produce").Inc()

	partition := s.resolvePartition(topic, req.Partition, req.Key)

	offset, err := s.produceWithRetry(ctx, topic, partition, req)
	if err != nil {
		if isLeadershipErr(err) {
			metrics.ProduceErrors.WithLabelValues(topic, "not_leader").Inc()
			leaderAddr, _ := s.broker.GetPartitionLeader(topic, partition)
			return nil, status.Errorf(codes.Unavailable, "not partition leader, try: %s", leaderAddr)
		}
		if errors.Is(err, broker.ErrTopicNotFound) {
			return nil, status.Errorf(codes.NotFound, "topic not found: %s", topic)
		}
		metrics.ProduceErrors.WithLabelValues(topic, "internal").Inc()
		return nil, status.Errorf(codes.Internal, "produce failed: %v", err)
	}

	metrics.MessagesProduced.WithLabelValues(topic).Inc()
	metrics.ProduceLatency.WithLabelValues(topic).Observe(time.Since(start).Seconds())

	return &api.ProduceResponse{Offset: offset, Partition: int32(partition)}, nil
}

// produceWithRetry writes locally when this broker leads the partition and
// proxies to the leader otherwise. During failover, leadership errors are
// retried within a bounded budget so clients survive broker failures without
// their own redirect logic. Requests that were already proxied once fail
// fast instead.
func (s *DistributedServer) produceWithRetry(ctx context.Context, topic string, partition int, req *api.ProduceRequest) (uint64, error) {
	deadline := time.Now().Add(produceRetryBudget)

	for {
		offset, err := s.broker.Produce(ctx, topic, partition, req.Record)
		if err == nil {
			return offset, nil
		}
		if !isLeadershipErr(err) {
			return 0, err
		}
		if isForwarded(ctx) {
			// The origin broker owns the retry loop.
			return 0, err
		}

		if client, cerr := s.partitionLeaderClient(topic, partition); cerr == nil {
			fwdCtx := metadata.AppendToOutgoingContext(ctx, forwardedMDKey, "1")
			resp, perr := client.Produce(fwdCtx, &api.ProduceRequest{
				Topic:     topic,
				Record:    req.Record,
				Key:       req.Key,
				Partition: int32(partition),
			})
			if perr == nil {
				return resp.Offset, nil
			}
		}

		if time.Now().After(deadline) || ctx.Err() != nil {
			return 0, err
		}
		time.Sleep(500 * time.Millisecond)
	}
}

func (s *DistributedServer) ProduceBatch(ctx context.Context, req *api.ProduceBatchRequest) (*api.ProduceBatchResponse, error) {
	start := time.Now()
	topic := req.Topic
	if topic == "" {
		topic = "default"
	}

	metrics.GRPCRequestsTotal.WithLabelValues("ProduceBatch").Inc()

	records := make([][]byte, len(req.Records))
	var firstKey []byte
	for i, r := range req.Records {
		records[i] = r.Value
		if firstKey == nil && len(r.Key) > 0 {
			firstKey = r.Key
		}
	}

	partition := s.resolvePartition(topic, req.Partition, firstKey)

	baseOffset, count, err := s.produceBatchWithRetry(ctx, topic, partition, req, records)
	if err != nil {
		if isLeadershipErr(err) {
			metrics.ProduceErrors.WithLabelValues(topic, "not_leader").Inc()
			leaderAddr, _ := s.broker.GetPartitionLeader(topic, partition)
			return nil, status.Errorf(codes.Unavailable, "not partition leader, try: %s", leaderAddr)
		}
		metrics.ProduceErrors.WithLabelValues(topic, "internal").Inc()
		return nil, status.Errorf(codes.Internal, "produce batch failed: %v", err)
	}

	metrics.MessagesProduced.WithLabelValues(topic).Add(float64(count))
	metrics.ProduceLatency.WithLabelValues(topic).Observe(time.Since(start).Seconds())

	return &api.ProduceBatchResponse{
		BaseOffset:  baseOffset,
		RecordCount: int32(count),
		Partition:   int32(partition),
	}, nil
}

func (s *DistributedServer) produceBatchWithRetry(ctx context.Context, topic string, partition int, req *api.ProduceBatchRequest, records [][]byte) (uint64, int, error) {
	deadline := time.Now().Add(produceRetryBudget)

	for {
		baseOffset, count, err := s.broker.ProduceBatch(ctx, topic, partition, records)
		if err == nil {
			return baseOffset, count, nil
		}
		if !isLeadershipErr(err) {
			return 0, 0, err
		}
		if isForwarded(ctx) {
			return 0, 0, err
		}

		if client, cerr := s.partitionLeaderClient(topic, partition); cerr == nil {
			fwdCtx := metadata.AppendToOutgoingContext(ctx, forwardedMDKey, "1")
			resp, perr := client.ProduceBatch(fwdCtx, &api.ProduceBatchRequest{
				Topic:     topic,
				Records:   req.Records,
				Partition: int32(partition),
			})
			if perr == nil {
				return resp.BaseOffset, int(resp.RecordCount), nil
			}
		}

		if time.Now().After(deadline) || ctx.Err() != nil {
			return 0, 0, err
		}
		time.Sleep(500 * time.Millisecond)
	}
}

func (s *DistributedServer) Consume(ctx context.Context, req *api.ConsumeRequest) (*api.ConsumeResponse, error) {
	start := time.Now()
	topic := req.Topic
	if topic == "" {
		topic = "default"
	}

	metrics.GRPCRequestsTotal.WithLabelValues("Consume").Inc()

	partition := int(req.Partition)
	if partition < 0 {
		partition = 0
	}

	record, err := s.broker.Consume(topic, partition, req.Offset)
	if err != nil {
		// Not hosted here: proxy the read to the partition leader.
		if errors.Is(err, broker.ErrPartitionNotFound) && !isForwarded(ctx) {
			if client, cerr := s.partitionLeaderClient(topic, partition); cerr == nil {
				fwdCtx := metadata.AppendToOutgoingContext(ctx, forwardedMDKey, "1")
				if resp, perr := client.Consume(fwdCtx, req); perr == nil {
					return resp, nil
				}
			}
		}
		if errors.Is(err, broker.ErrTopicNotFound) {
			return nil, status.Errorf(codes.NotFound, "topic not found: %s", topic)
		}
		return nil, status.Errorf(codes.NotFound, "offset not available: %v", err)
	}

	metrics.MessagesConsumed.WithLabelValues(topic).Inc()
	metrics.ConsumeLatency.WithLabelValues(topic).Observe(time.Since(start).Seconds())

	return &api.ConsumeResponse{
		Record:    record,
		Offset:    req.Offset,
		Partition: int32(partition),
		Timestamp: time.Now().UnixMilli(),
	}, nil
}

func (s *DistributedServer) ConsumeBatch(ctx context.Context, req *api.ConsumeBatchRequest) (*api.ConsumeBatchResponse, error) {
	start := time.Now()
	topic := req.Topic
	if topic == "" {
		topic = "default"
	}

	metrics.GRPCRequestsTotal.WithLabelValues("ConsumeBatch").Inc()

	partition := int(req.Partition)
	if partition < 0 {
		partition = 0
	}
	maxRecords := int(req.MaxRecords)
	if maxRecords <= 0 {
		maxRecords = 100
	}

	records, nextOffset, err := s.broker.ConsumeBatch(topic, partition, req.Offset, maxRecords)
	if err != nil {
		if errors.Is(err, broker.ErrPartitionNotFound) && !isForwarded(ctx) {
			if client, cerr := s.partitionLeaderClient(topic, partition); cerr == nil {
				fwdCtx := metadata.AppendToOutgoingContext(ctx, forwardedMDKey, "1")
				if resp, perr := client.ConsumeBatch(fwdCtx, req); perr == nil {
					return resp, nil
				}
			}
		}
		if errors.Is(err, broker.ErrTopicNotFound) {
			return nil, status.Errorf(codes.NotFound, "topic not found: %s", topic)
		}
		return nil, status.Errorf(codes.Internal, "consume batch failed: %v", err)
	}

	responses := make([]*api.ConsumeResponse, len(records))
	now := time.Now().UnixMilli()
	for i, record := range records {
		responses[i] = &api.ConsumeResponse{
			Record:    record,
			Offset:    req.Offset + uint64(i),
			Partition: int32(partition),
			Timestamp: now,
		}
	}

	metrics.MessagesConsumed.WithLabelValues(topic).Add(float64(len(records)))
	metrics.ConsumeLatency.WithLabelValues(topic).Observe(time.Since(start).Seconds())

	return &api.ConsumeBatchResponse{Records: responses, NextOffset: nextOffset}, nil
}

func (s *DistributedServer) ConsumeStream(req *api.ConsumeRequest, stream api.Log_ConsumeStreamServer) error {
	topic := req.Topic
	if topic == "" {
		topic = "default"
	}

	metrics.GRPCRequestsTotal.WithLabelValues("ConsumeStream").Inc()

	offset := req.Offset
	partition := int(req.Partition)
	ctx := stream.Context()

	const (
		minPollInterval = 10 * time.Millisecond
		maxPollInterval = 1 * time.Second
	)
	pollInterval := minPollInterval

	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
		}

		record, err := s.broker.Consume(topic, partition, offset)
		if err != nil {
			select {
			case <-ctx.Done():
				return ctx.Err()
			case <-time.After(pollInterval):
				pollInterval = pollInterval * 2
				if pollInterval > maxPollInterval {
					pollInterval = maxPollInterval
				}
				continue
			}
		}

		pollInterval = minPollInterval

		if err := stream.Send(&api.ConsumeResponse{
			Record:    record,
			Offset:    offset,
			Partition: int32(partition),
			Timestamp: time.Now().UnixMilli(),
		}); err != nil {
			return err
		}

		metrics.MessagesConsumed.WithLabelValues(topic).Inc()
		offset++
	}
}

func (s *DistributedServer) JoinConsumerGroup(ctx context.Context, req *api.JoinConsumerGroupRequest) (*api.JoinConsumerGroupResponse, error) {
	metrics.GRPCRequestsTotal.WithLabelValues("JoinConsumerGroup").Inc()

	partitions, genID, err := s.broker.JoinConsumerGroup(req.GroupId, req.ConsumerId, req.Topic)
	if err != nil {
		if errors.Is(err, broker.ErrNotLeader) {
			return nil, status.Errorf(codes.Unavailable, "not controller leader, try: %s", s.broker.LeaderGRPCAddr())
		}
		return nil, status.Errorf(codes.Internal, "join group failed: %v", err)
	}

	int32Partitions := make([]int32, len(partitions))
	for i, p := range partitions {
		int32Partitions[i] = int32(p)
	}

	metrics.ConsumerGroupMembers.WithLabelValues(req.GroupId).Inc()
	return &api.JoinConsumerGroupResponse{
		GroupId:            req.GroupId,
		AssignedPartitions: int32Partitions,
		GenerationId:       int32(genID),
	}, nil
}

func (s *DistributedServer) LeaveConsumerGroup(ctx context.Context, req *api.LeaveConsumerGroupRequest) (*api.LeaveConsumerGroupResponse, error) {
	metrics.GRPCRequestsTotal.WithLabelValues("LeaveConsumerGroup").Inc()
	if err := s.broker.LeaveConsumerGroup(req.GroupId, req.ConsumerId); err != nil {
		if errors.Is(err, broker.ErrNotLeader) {
			return nil, status.Errorf(codes.Unavailable, "not controller leader, try: %s", s.broker.LeaderGRPCAddr())
		}
		return nil, status.Errorf(codes.Internal, "leave group failed: %v", err)
	}
	metrics.ConsumerGroupMembers.WithLabelValues(req.GroupId).Dec()
	return &api.LeaveConsumerGroupResponse{}, nil
}

func (s *DistributedServer) CommitOffset(ctx context.Context, req *api.CommitOffsetRequest) (*api.CommitOffsetResponse, error) {
	metrics.GRPCRequestsTotal.WithLabelValues("CommitOffset").Inc()
	if err := s.broker.CommitOffset(req.GroupId, req.ConsumerId, req.Topic, int(req.Partition), req.Offset); err != nil {
		if errors.Is(err, broker.ErrNotLeader) {
			return nil, status.Errorf(codes.Unavailable, "not controller leader, try: %s", s.broker.LeaderGRPCAddr())
		}
		return nil, status.Errorf(codes.Internal, "commit offset failed: %v", err)
	}
	metrics.OffsetCommits.WithLabelValues(req.GroupId, req.Topic).Inc()
	return &api.CommitOffsetResponse{}, nil
}

func (s *DistributedServer) FetchOffset(ctx context.Context, req *api.FetchOffsetRequest) (*api.FetchOffsetResponse, error) {
	metrics.GRPCRequestsTotal.WithLabelValues("FetchOffset").Inc()
	offset, err := s.broker.FetchOffset(req.GroupId, req.ConsumerId, req.Topic, int(req.Partition))
	if err != nil {
		return nil, status.Errorf(codes.NotFound, "fetch offset failed: %v", err)
	}
	return &api.FetchOffsetResponse{Offset: offset}, nil
}

func (s *DistributedServer) Heartbeat(ctx context.Context, req *api.HeartbeatRequest) (*api.HeartbeatResponse, error) {
	metrics.GRPCRequestsTotal.WithLabelValues("Heartbeat").Inc()

	rebalanceNeeded, err := s.broker.Heartbeat(req.GroupId, req.ConsumerId, int(req.GenerationId))
	if err != nil {
		if errors.Is(err, broker.ErrNotLeader) {
			return nil, status.Errorf(codes.Unavailable, "not controller leader, try: %s", s.broker.LeaderGRPCAddr())
		}
		return nil, status.Errorf(codes.NotFound, "heartbeat failed: %v", err)
	}
	return &api.HeartbeatResponse{RebalanceNeeded: rebalanceNeeded}, nil
}

func (s *DistributedServer) GetAssignment(ctx context.Context, req *api.GetAssignmentRequest) (*api.GetAssignmentResponse, error) {
	metrics.GRPCRequestsTotal.WithLabelValues("GetAssignment").Inc()
	partitions, genID, err := s.broker.GetAssignment(req.GroupId, req.ConsumerId)
	if err != nil {
		return nil, status.Errorf(codes.NotFound, "assignment not found: %v", err)
	}
	int32Partitions := make([]int32, len(partitions))
	for i, p := range partitions {
		int32Partitions[i] = int32(p)
	}
	return &api.GetAssignmentResponse{
		Partitions:   int32Partitions,
		GenerationId: int32(genID),
	}, nil
}

func (s *DistributedServer) CreateTopic(ctx context.Context, req *api.CreateTopicRequest) (*api.CreateTopicResponse, error) {
	metrics.GRPCRequestsTotal.WithLabelValues("CreateTopic").Inc()

	partitionCount := int(req.PartitionCount)
	if partitionCount <= 0 {
		partitionCount = 1
	}

	// Replication factor scales with the live cluster size, capped at 3.
	replicationFactor := 1
	liveBrokers := 0
	for _, b := range s.broker.GetBrokers() {
		if b.Alive {
			liveBrokers++
		}
	}
	if liveBrokers >= 3 {
		replicationFactor = 3
	} else if liveBrokers == 2 {
		replicationFactor = 2
	}

	if err := s.broker.CreateTopic(req.Name, partitionCount, replicationFactor); err != nil {
		if errors.Is(err, broker.ErrNotLeader) {
			return nil, status.Errorf(codes.Unavailable, "not controller leader, try: %s", s.broker.LeaderGRPCAddr())
		}
		return nil, status.Errorf(codes.Internal, "create topic failed: %v", err)
	}

	metrics.TopicCount.Inc()
	return &api.CreateTopicResponse{Name: req.Name, PartitionCount: int32(partitionCount)}, nil
}

func (s *DistributedServer) ListTopics(ctx context.Context, req *api.ListTopicsRequest) (*api.ListTopicsResponse, error) {
	metrics.GRPCRequestsTotal.WithLabelValues("ListTopics").Inc()
	return &api.ListTopicsResponse{Topics: s.broker.ListTopics()}, nil
}

func (s *DistributedServer) GetTopicInfo(ctx context.Context, req *api.GetTopicInfoRequest) (*api.GetTopicInfoResponse, error) {
	metrics.GRPCRequestsTotal.WithLabelValues("GetTopicInfo").Inc()

	oldest, newest, partitionCount, err := s.broker.GetTopicInfo(req.Topic)
	if err != nil {
		return nil, status.Errorf(codes.NotFound, "topic not found: %s", req.Topic)
	}

	// Per-partition ranges are filled from local replicas where this broker
	// hosts one; remote-only partitions report zero ranges.
	partitions := make([]*api.PartitionInfo, partitionCount)
	for i := 0; i < partitionCount; i++ {
		pOldest, pNewest, ok := s.broker.GetPartitionOffsets(req.Topic, i)
		if !ok {
			pOldest, pNewest = 0, 0
		}
		partitions[i] = &api.PartitionInfo{
			PartitionId:  int32(i),
			OldestOffset: pOldest,
			NewestOffset: pNewest,
		}
	}

	metrics.TopicMessages.WithLabelValues(req.Topic).Set(float64(newest - oldest))
	return &api.GetTopicInfoResponse{
		Name:           req.Topic,
		OldestOffset:   oldest,
		NewestOffset:   newest,
		PartitionCount: int32(partitionCount),
		Partitions:     partitions,
	}, nil
}

func (s *DistributedServer) SetRetentionPolicy(ctx context.Context, req *api.SetRetentionPolicyRequest) (*api.SetRetentionPolicyResponse, error) {
	metrics.GRPCRequestsTotal.WithLabelValues("SetRetentionPolicy").Inc()
	maxAge := time.Duration(req.MaxAgeSeconds) * time.Second
	s.broker.SetRetentionPolicy(req.Topic, maxAge, req.MaxBytes)
	return &api.SetRetentionPolicyResponse{}, nil
}

func (s *DistributedServer) DeleteTopic(ctx context.Context, req *api.DeleteTopicRequest) (*api.DeleteTopicResponse, error) {
	metrics.GRPCRequestsTotal.WithLabelValues("DeleteTopic").Inc()

	if err := s.broker.DeleteTopic(req.Name); err != nil {
		if errors.Is(err, broker.ErrNotLeader) {
			return nil, status.Errorf(codes.Unavailable, "not controller leader, try: %s", s.broker.LeaderGRPCAddr())
		}
		return nil, status.Errorf(codes.Internal, "delete topic failed: %v", err)
	}

	metrics.TopicCount.Dec()
	return &api.DeleteTopicResponse{}, nil
}

// Join adds a node to the Raft cluster and registers it as a broker so the
// controller can assign partitions to it.
func (s *DistributedServer) Join(ctx context.Context, req *api.JoinRequest) (*api.JoinResponse, error) {
	metrics.GRPCRequestsTotal.WithLabelValues("Join").Inc()

	if err := s.broker.Join(req.NodeId, req.Addr); err != nil {
		if errors.Is(err, broker.ErrNotLeader) {
			return nil, status.Error(codes.Unavailable, "not leader")
		}
		return nil, status.Errorf(codes.Internal, "join failed: %v", err)
	}

	// Older clients may omit the extra addresses; they join Raft but can't
	// host partitions until they register.
	if req.GrpcAddr != "" {
		if err := s.broker.RegisterBroker(req.NodeId, req.GrpcAddr, req.Addr, req.DataAddr); err != nil {
			return nil, status.Errorf(codes.Internal, "register broker: %v", err)
		}
	}

	return &api.JoinResponse{}, nil
}

func (s *DistributedServer) GetServers(ctx context.Context, req *api.GetServersRequest) (*api.GetServersResponse, error) {
	metrics.GRPCRequestsTotal.WithLabelValues("GetServers").Inc()

	state := 0
	if s.broker.IsLeader() {
		state = 2
	}
	metrics.RaftState.WithLabelValues(s.config.NodeID).Set(float64(state))

	leaderID := s.broker.LeaderID()
	brokers := s.broker.GetBrokers()
	servers := make([]*api.Server, 0, len(brokers))
	for _, b := range brokers {
		servers = append(servers, &api.Server{
			Id:       b.ID,
			Address:  b.Address,
			IsLeader: b.ID == leaderID,
		})
	}

	return &api.GetServersResponse{Servers: servers}, nil
}