ReferenceHash balancer (#906)

* ReferenceHash balancer

* ReferenceHash randomBalancer

Author: fracasula

README.md

@@ -468,9 +468,13 @@ The `Writer` will return an error if it detects this ambiguity.
 
 #### Sarama
 
-If you're switching from Sarama and need/want to use the same algorithm for message
-partitioning, you can use the ```kafka.Hash``` balancer.  ```kafka.Hash``` routes
-messages to the same partitions that Sarama's default partitioner would route to.
+If you're switching from Sarama and need/want to use the same algorithm for message partitioning, you can either use 
+the `kafka.Hash` balancer or the `kafka.ReferenceHash` balancer:
+* `kafka.Hash` = `sarama.NewHashPartitioner`
+* `kafka.ReferenceHash` = `sarama.NewReferenceHashPartitioner`
+
+The `kafka.Hash` and `kafka.ReferenceHash` balancers would route messages to the same partitions that the two 
+aforementioned Sarama partitioners would route them to.
 
 ```go
 w := &kafka.Writer{

balancer.go

@@ -167,6 +167,53 @@ func (h *Hash) Balance(msg Message, partitions ...int) int {
 	return int(partition)
 }
 
+// ReferenceHash 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:
+//
+//      (int32(hasher.Sum32()) & 0x7fffffff) % len(partitions) => partition
+//
+// By default, ReferenceHash uses the FNV-1a algorithm. This is the same algorithm as
+// the Sarama NewReferenceHashPartitioner and ensures that messages produced by kafka-go will
+// be delivered to the same topics that the Sarama producer would be delivered to.
+type ReferenceHash struct {
+	rr     randomBalancer
+	Hasher hash.Hash32
+
+	// lock protects Hasher while calculating the hash code.  It is assumed that
+	// the Hasher field is read-only once the Balancer is created, so as a
+	// performance optimization, reads of the field are not protected.
+	lock sync.Mutex
+}
+
+func (h *ReferenceHash) Balance(msg Message, partitions ...int) int {
+	if msg.Key == nil {
+		return h.rr.Balance(msg, partitions...)
+	}
+
+	hasher := h.Hasher
+	if hasher != nil {
+		h.lock.Lock()
+		defer h.lock.Unlock()
+	} else {
+		hasher = fnv1aPool.Get().(hash.Hash32)
+		defer fnv1aPool.Put(hasher)
+	}
+
+	hasher.Reset()
+	if _, err := hasher.Write(msg.Key); err != nil {
+		panic(err)
+	}
+
+	// uses the same algorithm as the Sarama's referenceHashPartitioner.
+	// note the type conversions here. if the uint32 hash code is not cast to
+	// an int32, we do not get the same result as sarama.
+	partition := (int32(hasher.Sum32()) & 0x7fffffff) % int32(len(partitions))
+	return int(partition)
+}
+
 type randomBalancer struct {
 	mock int // mocked return value, used for testing
 }

balancer_test.go

@@ -63,6 +63,65 @@ func TestHashBalancer(t *testing.T) {
 	}
 }
 
+func TestReferenceHashBalancer(t *testing.T) {
+	testCases := map[string]struct {
+		Key               []byte
+		Hasher            hash.Hash32
+		Partitions        []int
+		Partition         int
+		RndBalancerResult int
+	}{
+		"nil": {
+			Key:               nil, // nil key means random partition
+			Partitions:        []int{0, 1, 2},
+			Partition:         123,
+			RndBalancerResult: 123,
+		},
+		"partition-0": {
+			Key:        []byte("blah"),
+			Partitions: []int{0, 1},
+			Partition:  0,
+		},
+		"partition-1": {
+			Key:        []byte("blah"),
+			Partitions: []int{0, 1, 2},
+			Partition:  1,
+		},
+		"partition-2": {
+			Key:        []byte("castle"),
+			Partitions: []int{0, 1, 2},
+			Partition:  2,
+		},
+		"custom hash": {
+			Key:        []byte("boop"),
+			Hasher:     crc32.NewIEEE(),
+			Partitions: []int{0, 1, 2},
+			Partition:  1,
+		},
+		"hash code with MSB set": {
+			Key:        []byte("20"),
+			Partitions: []int{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15},
+			Partition:  15,
+		},
+	}
+
+	for label, test := range testCases {
+		t.Run(label, func(t *testing.T) {
+			var rr randomBalancer
+			if test.Key == nil {
+				rr.mock = test.RndBalancerResult
+			}
+
+			msg := Message{Key: test.Key}
+			h := ReferenceHash{Hasher: test.Hasher, rr: rr}
+			partition := h.Balance(msg, test.Partitions...)
+			if partition != test.Partition {
+				t.Errorf("expected %v; got %v", test.Partition, partition)
+			}
+		})
+	}
+}
+
 func TestCRC32Balancer(t *testing.T) {
 	// These tests are taken from the default "consistent_random" partitioner from
 	// https://github.com/edenhill/librdkafka/blob/master/tests/0048-partitioner.c