func (p *BalancePlan) AdjustCooperative(b *ConsumerBalancer) {
	allAdded := make(map[string]map[int32]string, 100)  
	allRevoked := make(map[string]map[int32]struct{}, 100)
	addT := func(t string) map[int32]string {
		addT := allAdded[t]
		if addT == nil {
			addT = make(map[int32]string, 20)
			allAdded[t] = addT
		}
		return addT
	}
	revokeT := func(t string) map[int32]struct{} {
		revokeT := allRevoked[t]
		if revokeT == nil {
			revokeT = make(map[int32]struct{}, 20)
			allRevoked[t] = revokeT
		}
		return revokeT
	}
	tmap := make(map[string]struct{})  
	pmap := make(map[int32]struct{})   
	plan := p.plan
	b.EachMember(func(member *kmsg.JoinGroupResponseMember, meta *kmsg.GroupMemberMetadata) {
		planned := plan[member.MemberID]
		for ptopic := range planned {  
			tmap[ptopic] = struct{}{}
		}
		for _, otopic := range meta.OwnedPartitions {  
			topic := otopic.Topic
			delete(tmap, topic)
			ppartitions, exists := planned[topic]
			if !exists {  
				allRevokedT := revokeT(topic)
				for _, opartition := range otopic.Partitions {
					allRevokedT[opartition] = struct{}{}
				}
				continue
			}
			for _, ppartition := range ppartitions {
				pmap[ppartition] = struct{}{}
			}
			for _, opartition := range otopic.Partitions {
				delete(pmap, opartition)
			}
			if len(pmap) > 0 {
				allAddedT := addT(topic)
				for ppartition := range pmap {
					delete(pmap, ppartition)
					allAddedT[ppartition] = member.MemberID
				}
			}
			for _, opartition := range otopic.Partitions {
				pmap[opartition] = struct{}{}
			}
			for _, ppartition := range ppartitions {
				delete(pmap, ppartition)
			}
			if len(pmap) > 0 {
				allRevokedT := revokeT(topic)
				for opartition := range pmap {
					delete(pmap, opartition)
					allRevokedT[opartition] = struct{}{}
				}
			}
		}
		for ptopic := range tmap {  
			delete(tmap, ptopic)
			allAddedT := addT(ptopic)
			for _, ppartition := range planned[ptopic] {
				allAddedT[ppartition] = member.MemberID
			}
		}
	})
	for topic, rpartitions := range allRevoked {
		atopic, exists := allAdded[topic]
		if !exists {
			continue
		}
		for rpartition := range rpartitions {
			amember, exists := atopic[rpartition]
			if !exists {
				continue
			}
			ptopics := plan[amember]
			ppartitions := ptopics[topic]
			for i, ppartition := range ppartitions {
				if ppartition == rpartition {
					ppartitions[i] = ppartitions[len(ppartitions)-1]
					ppartitions = ppartitions[:len(ppartitions)-1]
					break
				}
			}
			if len(ppartitions) > 0 {
				ptopics[topic] = ppartitions
			} else {
				delete(ptopics, topic)
			}
		}
	}
}
AdjustCooperative performs the final adjustment to a plan for cooperative
balancing.

Over the plan, we remove all partitions that migrated from one member (where
it was assigned) to a new member (where it is now planned).

This allows members that had partitions removed to revoke and rejoin, which
will then do another rebalance, and in that new rebalance, the planned
partitions are now on the free list to be assigned.
package kgo
import (
	"bytes"
	"fmt"
	"sort"
	"strings"
	"github.com/twmb/franz-go/pkg/kerr"
	"github.com/twmb/franz-go/pkg/kgo/internal/sticky"
	"github.com/twmb/franz-go/pkg/kmsg"
)
type GroupBalancer interface {
	ProtocolName() string
	JoinGroupMetadata(
		topicInterests []string,
		currentAssignment map[string][]int32,
		generation int32,
	) []byte
	ParseSyncAssignment(assignment []byte) (map[string][]int32, error)
	MemberBalancer(members []kmsg.JoinGroupResponseMember) (b GroupMemberBalancer, topics map[string]struct{}, err error)
	IsCooperative() bool
}
type GroupMemberBalancer interface {
	Balance(topics map[string]int32) IntoSyncAssignment
}
type IntoSyncAssignment interface {
	IntoSyncAssignment() []kmsg.SyncGroupRequestGroupAssignment
}
type ConsumerBalancer struct {
	b         ConsumerBalancerBalance
	members   []kmsg.JoinGroupResponseMember
	metadatas []kmsg.GroupMemberMetadata
	topics    map[string]struct{}
}
func (b *ConsumerBalancer) Balance(topics map[string]int32) IntoSyncAssignment {
	return b.b.Balance(b, topics)
}
func (b *ConsumerBalancer) Members() []kmsg.JoinGroupResponseMember {
	return b.members
}
func (b *ConsumerBalancer) EachMember(fn func(member *kmsg.JoinGroupResponseMember, meta *kmsg.GroupMemberMetadata)) {
	for i := range b.members {
		fn(&b.members[i], &b.metadatas[i])
	}
}
func (b *ConsumerBalancer) MemberAt(n int) (*kmsg.JoinGroupResponseMember, *kmsg.GroupMemberMetadata) {
	return &b.members[n], &b.metadatas[n]
}
func (b *ConsumerBalancer) MemberTopics() map[string]struct{} {
	if b == nil {
		return nil
	}
	return b.topics
}
func (b *ConsumerBalancer) NewPlan() *BalancePlan {
	plan := make(map[string]map[string][]int32, len(b.members))
	for i := range b.members {
		plan[b.members[i].MemberID] = make(map[string][]int32)
	}
	return &BalancePlan{plan}
}
type ConsumerBalancerBalance interface {
	Balance(*ConsumerBalancer, map[string]int32) IntoSyncAssignment
}
func ParseConsumerSyncAssignment(assignment []byte) (map[string][]int32, error) {
	var kassignment kmsg.GroupMemberAssignment
	if err := kassignment.ReadFrom(assignment); err != nil {
		return nil, fmt.Errorf("sync assignment parse failed: %v", err)
	}
	m := make(map[string][]int32, len(kassignment.Topics))
	for _, topic := range kassignment.Topics {
		m[topic.Topic] = topic.Partitions
	}
	return m, nil
}
func NewConsumerBalancer(balance ConsumerBalancerBalance, members []kmsg.JoinGroupResponseMember) (*ConsumerBalancer, error) {
	b := &ConsumerBalancer{
		b:         balance,
		members:   members,
		metadatas: make([]kmsg.GroupMemberMetadata, len(members)),
		topics:    make(map[string]struct{}),
	}
	for i, member := range members {
		meta := &b.metadatas[i]
		meta.Default()
		if err := meta.ReadFrom(member.ProtocolMetadata); err != nil {
			return nil, fmt.Errorf("unable to read member metadata: %v", err)
		}
		for _, topic := range meta.Topics {
			b.topics[topic] = struct{}{}
		}
		sort.Strings(meta.Topics)
	}
	return b, nil
}
type BalancePlan struct {
	plan map[string]map[string][]int32  
}
func (p *BalancePlan) String() string {
	var sb strings.Builder
	var membersWritten int
	for member, topics := range p.plan {
		membersWritten++
		sb.WriteString(member)
		sb.WriteString("{")
		var topicsWritten int
		for topic, partitions := range topics {
			fmt.Fprintf(&sb, "%s%v", topic, partitions)
			topicsWritten++
			if topicsWritten < len(topics) {
				sb.WriteString(", ")
			}
		}
		sb.WriteString("}")
		if membersWritten < len(p.plan) {
			sb.WriteString(", ")
		}
	}
	return sb.String()
}
func (p *BalancePlan) AddPartition(member *kmsg.JoinGroupResponseMember, topic string, partition int32) {
	memberPlan := p.plan[member.MemberID]
	memberPlan[topic] = append(memberPlan[topic], partition)
}
func (p *BalancePlan) AddPartitions(member *kmsg.JoinGroupResponseMember, topic string, partitions []int32) {
	memberPlan := p.plan[member.MemberID]
	memberPlan[topic] = append(memberPlan[topic], partitions...)
}
func (p *BalancePlan) IntoSyncAssignment() []kmsg.SyncGroupRequestGroupAssignment {
	kassignments := make([]kmsg.SyncGroupRequestGroupAssignment, 0, len(p.plan))
	for member, assignment := range p.plan {
		var kassignment kmsg.GroupMemberAssignment
		for topic, partitions := range assignment {
			sort.Slice(partitions, func(i, j int) bool { return partitions[i] < partitions[j] })
			assnTopic := kmsg.NewGroupMemberAssignmentTopic()
			assnTopic.Topic = topic
			assnTopic.Partitions = partitions
			kassignment.Topics = append(kassignment.Topics, assnTopic)
		}
		sort.Slice(kassignment.Topics, func(i, j int) bool { return kassignment.Topics[i].Topic < kassignment.Topics[j].Topic })
		syncAssn := kmsg.NewSyncGroupRequestGroupAssignment()
		syncAssn.MemberID = member
		syncAssn.MemberAssignment = kassignment.AppendTo(nil)
		kassignments = append(kassignments, syncAssn)
	}
	sort.Slice(kassignments, func(i, j int) bool { return kassignments[i].MemberID < kassignments[j].MemberID })
	return kassignments
}
func joinMemberLess(l, r *kmsg.JoinGroupResponseMember) bool {
	if l.InstanceID != nil {
		if r.InstanceID == nil {
			return true
		}
		return *l.InstanceID < *r.InstanceID
	}
	if r.InstanceID != nil {
		return false
	}
	return l.MemberID < r.MemberID
}
func sortJoinMembers(members []kmsg.JoinGroupResponseMember) {
	sort.Slice(members, func(i, j int) bool { return joinMemberLess(&members[i], &members[j]) })
}
func sortJoinMemberPtrs(members []*kmsg.JoinGroupResponseMember) {
	sort.Slice(members, func(i, j int) bool { return joinMemberLess(members[i], members[j]) })
}
func (g *groupConsumer) findBalancer(from, proto string) (GroupBalancer, error) {
	for _, b := range g.cfg.balancers {
		if b.ProtocolName() == proto {
			return b, nil
		}
	}
	var ours []string
	for _, b := range g.cfg.balancers {
		ours = append(ours, b.ProtocolName())
	}
	g.cl.cfg.logger.Log(LogLevelError, fmt.Sprintf("%s could not find Kafka-chosen balancer", from), "kafka_choice", proto, "our_set", strings.Join(ours, ", "))
	return nil, fmt.Errorf("unable to balance: none of our balancers have a name equal to the balancer chosen for balancing (%s)", proto)
}
func (g *groupConsumer) balanceGroup(proto string, members []kmsg.JoinGroupResponseMember) ([]kmsg.SyncGroupRequestGroupAssignment, error) {
	g.cl.cfg.logger.Log(LogLevelInfo, "balancing group as leader")
	b, err := g.findBalancer("balance group", proto)
	if err != nil {
		return nil, err
	}
	sortJoinMembers(members)
	memberBalancer, topics, err := b.MemberBalancer(members)
	if err != nil {
		return nil, fmt.Errorf("unable to create group member balancer: %v", err)
	}
	myTopics := g.tps.load()
	var needMeta bool
	topicPartitionCount := make(map[string]int32, len(topics))
	for topic := range topics {
		data, exists := myTopics[topic]
		if !exists {
			needMeta = true
			continue
		}
		topicPartitionCount[topic] = int32(len(data.load().partitions))
	}
	if needMeta {
		g.cl.cfg.logger.Log(LogLevelInfo, "group members indicated interest in topics the leader is not assigned, fetching metadata for all group topics")
		var metaTopics []string
		for topic := range topics {
			metaTopics = append(metaTopics, topic)
		}
		_, resp, err := g.cl.fetchMetadataForTopics(g.ctx, false, metaTopics)
		if err != nil {
			return nil, fmt.Errorf("unable to fetch metadata for group topics: %v", err)
		}
		for i := range resp.Topics {
			t := &resp.Topics[i]
			if t.Topic == nil {
				g.cl.cfg.logger.Log(LogLevelWarn, "metadata resp in balance for topic has nil topic, skipping...", "err", kerr.ErrorForCode(t.ErrorCode))
				continue
			}
			if t.ErrorCode != 0 {
				g.cl.cfg.logger.Log(LogLevelWarn, "metadata resp in balance for topic has error, skipping...", "topic", t.Topic, "err", kerr.ErrorForCode(t.ErrorCode))
				continue
			}
			topicPartitionCount[*t.Topic] = int32(len(t.Partitions))
		}
	}
	if b, ok := memberBalancer.(*ConsumerBalancer); ok {
		interests := new(bytes.Buffer)
		b.EachMember(func(member *kmsg.JoinGroupResponseMember, meta *kmsg.GroupMemberMetadata) {
			interests.Reset()
			fmt.Fprintf(interests, "interested topics: %v, previously owned: ", meta.Topics)
			for _, owned := range meta.OwnedPartitions {
				sort.Slice(owned.Partitions, func(i, j int) bool { return owned.Partitions[i] < owned.Partitions[j] })
				fmt.Fprintf(interests, "%s%v, ", owned.Topic, owned.Partitions)
			}
			strInterests := interests.String()
			strInterests = strings.TrimSuffix(strInterests, ", ")
			if member.InstanceID == nil {
				g.cl.cfg.logger.Log(LogLevelInfo, "balance group member", "id", member.MemberID, "interests", strInterests)
			} else {
				g.cl.cfg.logger.Log(LogLevelInfo, "balance group member", "id", member.MemberID, "instance_id", *member.InstanceID, "interests", strInterests)
			}
		})
	} else {
		g.cl.cfg.logger.Log(LogLevelInfo, "unable to log information about group member interests: the user has defined a custom balancer (not a *ConsumerBalancer)")
	}
	into := memberBalancer.Balance(topicPartitionCount)
	if p, ok := into.(*BalancePlan); ok {
		g.cl.cfg.logger.Log(LogLevelInfo, "balanced", "plan", p.String())
	} else {
		g.cl.cfg.logger.Log(LogLevelInfo, "unable to log balance plan: the user has returned a custom IntoSyncAssignment (not a *BalancePlan)")
	}
	return into.IntoSyncAssignment(), nil
}
func memberMetadataV0(interests []string) []byte {
	meta := kmsg.NewGroupMemberMetadata()
	meta.Version = 0
	meta.Topics = interests  
	return meta.AppendTo(nil)
}
func RoundRobinBalancer() GroupBalancer {
	return new(roundRobinBalancer)
}
type roundRobinBalancer struct{}
func (*roundRobinBalancer) ProtocolName() string { return "roundrobin" }
func (*roundRobinBalancer) IsCooperative() bool  { return false }
func (*roundRobinBalancer) JoinGroupMetadata(interests []string, _ map[string][]int32, _ int32) []byte {
	return memberMetadataV0(interests)
}
func (*roundRobinBalancer) ParseSyncAssignment(assignment []byte) (map[string][]int32, error) {
	return ParseConsumerSyncAssignment(assignment)
}
func (r *roundRobinBalancer) MemberBalancer(members []kmsg.JoinGroupResponseMember) (GroupMemberBalancer, map[string]struct{}, error) {
	b, err := NewConsumerBalancer(r, members)
	return b, b.MemberTopics(), err
}
func (*roundRobinBalancer) Balance(b *ConsumerBalancer, topics map[string]int32) IntoSyncAssignment {
	type topicPartition struct {
		topic     string
		partition int32
	}
	var nparts int
	for _, partitions := range topics {
		nparts += int(partitions)
	}
	allParts := make([]topicPartition, 0, nparts)
	for topic := range b.MemberTopics() {
		for partition := int32(0); partition < topics[topic]; partition++ {
			allParts = append(allParts, topicPartition{
				topic,
				partition,
			})
		}
	}
	sort.Slice(allParts, func(i, j int) bool {
		l, r := allParts[i], allParts[j]
		return l.topic < r.topic || l.topic == r.topic && l.partition < r.partition
	})
	plan := b.NewPlan()
	var memberIdx int
	for len(allParts) > 0 {
		next := allParts[0]
		allParts = allParts[1:]
	assigned:
		for {
			member, meta := b.MemberAt(memberIdx)
			memberIdx = (memberIdx + 1) % len(b.Members())
			for _, topic := range meta.Topics {
				if topic == next.topic {
					plan.AddPartition(member, next.topic, next.partition)
					break assigned
				}
			}
		}
	}
	return plan
}
func RangeBalancer() GroupBalancer {
	return new(rangeBalancer)
}
type rangeBalancer struct{}
func (*rangeBalancer) ProtocolName() string { return "range" }
func (*rangeBalancer) IsCooperative() bool  { return false }
func (*rangeBalancer) JoinGroupMetadata(interests []string, _ map[string][]int32, _ int32) []byte {
	return memberMetadataV0(interests)
}
func (*rangeBalancer) ParseSyncAssignment(assignment []byte) (map[string][]int32, error) {
	return ParseConsumerSyncAssignment(assignment)
}
func (r *rangeBalancer) MemberBalancer(members []kmsg.JoinGroupResponseMember) (GroupMemberBalancer, map[string]struct{}, error) {
	b, err := NewConsumerBalancer(r, members)
	return b, b.MemberTopics(), err
}
func (*rangeBalancer) Balance(b *ConsumerBalancer, topics map[string]int32) IntoSyncAssignment {
	topics2PotentialConsumers := make(map[string][]*kmsg.JoinGroupResponseMember)
	b.EachMember(func(member *kmsg.JoinGroupResponseMember, meta *kmsg.GroupMemberMetadata) {
		for _, topic := range meta.Topics {
			topics2PotentialConsumers[topic] = append(topics2PotentialConsumers[topic], member)
		}
	})
	plan := b.NewPlan()
	for topic, potentialConsumers := range topics2PotentialConsumers {
		sortJoinMemberPtrs(potentialConsumers)
		numPartitions := topics[topic]
		partitions := make([]int32, numPartitions)
		for i := range partitions {
			partitions[i] = int32(i)
		}
		numParts := len(partitions)
		div, rem := numParts/len(potentialConsumers), numParts%len(potentialConsumers)
		var consumerIdx int
		for len(partitions) > 0 {
			num := div
			if rem > 0 {
				num++
				rem--
			}
			member := potentialConsumers[consumerIdx]
			plan.AddPartitions(member, topic, partitions[:num])
			consumerIdx++
			partitions = partitions[num:]
		}
	}
	return plan
}
func StickyBalancer() GroupBalancer {
	return &stickyBalancer{cooperative: false}
}
type stickyBalancer struct {
	cooperative bool
}
func (s *stickyBalancer) ProtocolName() string {
	if s.cooperative {
		return "cooperative-sticky"
	}
	return "sticky"
}
func (s *stickyBalancer) IsCooperative() bool { return s.cooperative }
func (s *stickyBalancer) JoinGroupMetadata(interests []string, currentAssignment map[string][]int32, generation int32) []byte {
	meta := kmsg.NewGroupMemberMetadata()
	meta.Version = 0
	meta.Topics = interests
	if s.cooperative {
		meta.Version = 1
	}
	stickyMeta := kmsg.NewStickyMemberMetadata()
	stickyMeta.Generation = generation
	for topic, partitions := range currentAssignment {
		if s.cooperative {
			metaPart := kmsg.NewGroupMemberMetadataOwnedPartition()
			metaPart.Topic = topic
			metaPart.Partitions = partitions
			meta.OwnedPartitions = append(meta.OwnedPartitions, metaPart)
		}
		stickyAssn := kmsg.NewStickyMemberMetadataCurrentAssignment()
		stickyAssn.Topic = topic
		stickyAssn.Partitions = partitions
		stickyMeta.CurrentAssignment = append(stickyMeta.CurrentAssignment, stickyAssn)
	}
	metaOwned := meta.OwnedPartitions
	stickyCurrent := stickyMeta.CurrentAssignment
	sort.Slice(metaOwned, func(i, j int) bool { return metaOwned[i].Topic < metaOwned[j].Topic })
	sort.Slice(stickyCurrent, func(i, j int) bool { return stickyCurrent[i].Topic < stickyCurrent[j].Topic })
	meta.UserData = stickyMeta.AppendTo(nil)
	return meta.AppendTo(nil)
}
func (*stickyBalancer) ParseSyncAssignment(assignment []byte) (map[string][]int32, error) {
	return ParseConsumerSyncAssignment(assignment)
}
func (s *stickyBalancer) MemberBalancer(members []kmsg.JoinGroupResponseMember) (GroupMemberBalancer, map[string]struct{}, error) {
	b, err := NewConsumerBalancer(s, members)
	return b, b.MemberTopics(), err
}
func (s *stickyBalancer) Balance(b *ConsumerBalancer, topics map[string]int32) IntoSyncAssignment {
	stickyMembers := make([]sticky.GroupMember, 0, len(b.Members()))
	b.EachMember(func(member *kmsg.JoinGroupResponseMember, meta *kmsg.GroupMemberMetadata) {
		stickyMembers = append(stickyMembers, sticky.GroupMember{
			ID:       member.MemberID,
			Topics:   meta.Topics,
			UserData: meta.UserData,
		})
	})
	p := &BalancePlan{sticky.Balance(stickyMembers, topics)}
	if s.cooperative {
		p.AdjustCooperative(b)
	}
	return p
}
func CooperativeStickyBalancer() GroupBalancer {
	return &stickyBalancer{cooperative: true}
}