github.com/johnnyeven/libtools@v0.0.0-20191126065708-61829c1adf46/kafka/consumergroup/utils.go (about)

     1  package consumergroup
     2  
     3  import (
     4  	"crypto/rand"
     5  	"fmt"
     6  	"io"
     7  	"os"
     8  	"sort"
     9  
    10  	"github.com/wvanbergen/kazoo-go"
    11  )
    12  
    13  func retrievePartitionLeaders(partitions kazoo.PartitionList) (partitionLeaders, error) {
    14  
    15  	pls := make(partitionLeaders, 0, len(partitions))
    16  	for _, partition := range partitions {
    17  		leader, err := partition.Leader()
    18  		if err != nil {
    19  			return nil, err
    20  		}
    21  
    22  		pl := partitionLeader{id: partition.ID, leader: leader, partition: partition}
    23  		pls = append(pls, pl)
    24  	}
    25  
    26  	return pls, nil
    27  }
    28  
    29  // Divides a set of partitions between a set of consumers.
    30  func dividePartitionsBetweenConsumers(consumers kazoo.ConsumergroupInstanceList, partitions partitionLeaders) map[string][]*kazoo.Partition {
    31  	result := make(map[string][]*kazoo.Partition)
    32  
    33  	plen := len(partitions)
    34  	clen := len(consumers)
    35  	if clen == 0 {
    36  		return result
    37  	}
    38  
    39  	sort.Sort(partitions)
    40  	sort.Sort(consumers)
    41  
    42  	n := plen / clen
    43  	m := plen % clen
    44  	p := 0
    45  	for i, consumer := range consumers {
    46  		first := p
    47  		last := first + n
    48  		if m > 0 && i < m {
    49  			last++
    50  		}
    51  		if last > plen {
    52  			last = plen
    53  		}
    54  
    55  		for _, pl := range partitions[first:last] {
    56  			result[consumer.ID] = append(result[consumer.ID], pl.partition)
    57  		}
    58  		p = last
    59  	}
    60  
    61  	return result
    62  }
    63  
    64  type partitionLeader struct {
    65  	id        int32
    66  	leader    int32
    67  	partition *kazoo.Partition
    68  }
    69  
    70  // A sortable slice of PartitionLeader structs
    71  type partitionLeaders []partitionLeader
    72  
    73  func (pls partitionLeaders) Len() int {
    74  	return len(pls)
    75  }
    76  
    77  func (pls partitionLeaders) Less(i, j int) bool {
    78  	return pls[i].leader < pls[j].leader || (pls[i].leader == pls[j].leader && pls[i].id < pls[j].id)
    79  }
    80  
    81  func (s partitionLeaders) Swap(i, j int) {
    82  	s[i], s[j] = s[j], s[i]
    83  }
    84  
    85  func generateUUID() (string, error) {
    86  	uuid := make([]byte, 16)
    87  	n, err := io.ReadFull(rand.Reader, uuid)
    88  	if n != len(uuid) || err != nil {
    89  		return "", err
    90  	}
    91  	// variant bits; see section 4.1.1
    92  	uuid[8] = uuid[8]&^0xc0 | 0x80
    93  	// version 4 (pseudo-random); see section 4.1.3
    94  	uuid[6] = uuid[6]&^0xf0 | 0x40
    95  	return fmt.Sprintf("%x-%x-%x-%x-%x", uuid[0:4], uuid[4:6], uuid[6:8], uuid[8:10], uuid[10:]), nil
    96  }
    97  
    98  func generateConsumerID() (consumerID string, err error) {
    99  	var uuid, hostname string
   100  
   101  	uuid, err = generateUUID()
   102  	if err != nil {
   103  		return
   104  	}
   105  
   106  	hostname, err = os.Hostname()
   107  	if err != nil {
   108  		return
   109  	}
   110  
   111  	consumerID = fmt.Sprintf("%s:%s", hostname, uuid)
   112  	return
   113  }