1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
|
package kafka
import (
"hash"
"hash/fnv"
"sort"
"sync"
)
// The Balancer interface provides an abstraction of the message distribution
// logic used by Writer instances to route messages to the partitions available
// on a kafka cluster.
//
// Instances of Balancer do not have to be safe to use concurrently by multiple
// goroutines, the Writer implementation ensures that calls to Balance are
// synchronized.
type Balancer interface {
// Balance receives a message and a set of available partitions and
// returns the partition number that the message should be routed to.
//
// An application should refrain from using a balancer to manage multiple
// sets of partitions (from different topics for examples), use one balancer
// instance for each partition set, so the balancer can detect when the
// partitions change and assume that the kafka topic has been rebalanced.
Balance(msg Message, partitions ...int) (partition int)
}
// BalancerFunc is an implementation of the Balancer interface that makes it
// possible to use regular functions to distribute messages across partitions.
type BalancerFunc func(Message, ...int) int
// Balance calls f, satisfies the Balancer interface.
func (f BalancerFunc) Balance(msg Message, partitions ...int) int {
return f(msg, partitions...)
}
// RoundRobin is an Balancer implementation that equally distributes messages
// across all available partitions.
type RoundRobin struct {
offset uint64
}
// Balance satisfies the Balancer interface.
func (rr *RoundRobin) Balance(msg Message, partitions ...int) int {
length := uint64(len(partitions))
offset := rr.offset
rr.offset++
return partitions[offset%length]
}
// LeastBytes is a Balancer implementation that routes messages to the partition
// that has received the least amount of data.
//
// Note that no coordination is done between multiple producers, having good
// balancing relies on the fact that each producer using a LeastBytes balancer
// should produce well balanced messages.
type LeastBytes struct {
counters []leastBytesCounter
}
type leastBytesCounter struct {
partition int
bytes uint64
}
// Balance satisfies the Balancer interface.
func (lb *LeastBytes) Balance(msg Message, partitions ...int) int {
for _, p := range partitions {
if c := lb.counterOf(p); c == nil {
lb.counters = lb.makeCounters(partitions...)
break
}
}
minBytes := lb.counters[0].bytes
minIndex := 0
for i, c := range lb.counters[1:] {
if c.bytes < minBytes {
minIndex = i + 1
minBytes = c.bytes
}
}
c := &lb.counters[minIndex]
c.bytes += uint64(len(msg.Key)) + uint64(len(msg.Value))
return c.partition
}
func (lb *LeastBytes) counterOf(partition int) *leastBytesCounter {
i := sort.Search(len(lb.counters), func(i int) bool {
return lb.counters[i].partition >= partition
})
if i == len(lb.counters) || lb.counters[i].partition != partition {
return nil
}
return &lb.counters[i]
}
func (lb *LeastBytes) makeCounters(partitions ...int) (counters []leastBytesCounter) {
counters = make([]leastBytesCounter, len(partitions))
for i, p := range partitions {
counters[i].partition = p
}
sort.Slice(counters, func(i int, j int) bool {
return counters[i].partition < counters[j].partition
})
return
}
var (
fnv1aPool = &sync.Pool{
New: func() interface{} {
return fnv.New32a()
},
}
)
// Hash is a Balancer that uses the provided hash function to determine which
// partition to route messages to. This ensures that messages with the same key
// are routed to the same partition.
//
// The logic to calculate the partition is:
//
// hasher.Sum32() % len(partitions) => partition
//
// By default, Hash uses the FNV-1a algorithm. This is the same algorithm used
// by the Sarama Producer and ensures that messages produced by kafka-go will
// be delivered to the same topics that the Sarama producer would be delivered to
type Hash struct {
rr RoundRobin
Hasher hash.Hash32
}
func (h *Hash) Balance(msg Message, partitions ...int) (partition int) {
if msg.Key == nil {
return h.rr.Balance(msg, partitions...)
}
hasher := h.Hasher
if hasher == nil {
hasher = fnv1aPool.Get().(hash.Hash32)
defer fnv1aPool.Put(hasher)
}
hasher.Reset()
if _, err := hasher.Write(msg.Key); err != nil {
panic(err)
}
// uses same algorithm that Sarama's hashPartitioner uses
partition = int(hasher.Sum32()) % len(partitions)
if partition < 0 {
partition = -partition
}
return
}
|
