[Issue 172] Add key based batcher (#400)

Fixes #172

### Motivation

Add a new batch message container named `keyBasedBatchContainer` to support batching message in key_shared subscription mode.

### Modifications

- add `BatchBuilder` interface, add `FlushBatches` and `IsMultiBatches` func
- change old `BatchBuilder` struct to `batchContainer`
- add `keyBasedBatchContainer`
- add tests

### Verifying this change

This change added tests and can be verified as follows:

  - *Added integration tests for key based batch producer with multiple consumer in KeyShared mode*
  - *Added integration tests for message ordering with key based batch producer and KeyShared consumer*

go.mod

@@ -14,6 +14,8 @@ require (
 	github.com/klauspost/compress v1.10.8
 	github.com/kr/pretty v0.2.0 // indirect
 	github.com/linkedin/goavro/v2 v2.9.8
+	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
+	github.com/modern-go/reflect2 v1.0.1 // indirect
 	github.com/pierrec/lz4 v2.0.5+incompatible
 	github.com/pkg/errors v0.9.1
 	github.com/prometheus/client_golang v1.7.1

pulsar/batcher_builder.go

@@ -0,0 +1,44 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package pulsar
+
+import (
+	"errors"
+
+	"github.com/apache/pulsar-client-go/pulsar/internal"
+)
+
+type BatcherBuilderType int
+
+const (
+	DefaultBatchBuilder BatcherBuilderType = iota
+	KeyBasedBatchBuilder
+)
+
+func GetBatcherBuilderProvider(typ BatcherBuilderType) (
+	internal.BatcherBuilderProvider, error,
+) {
+	switch typ {
+	case DefaultBatchBuilder:
+		return internal.NewBatchBuilder, nil
+	case KeyBasedBatchBuilder:
+		return internal.NewKeyBasedBatchBuilder, nil
+	default:
+		return nil, errors.New("unsupported batcher builder provider type")
+	}
+}

pulsar/consumer_test.go

@@ -1712,3 +1712,170 @@ func TestConsumerName(t *testing.T) {
 
 	assert.Equal(consumerName, consumer.Name())
 }
+
+func TestKeyBasedBatchProducerConsumerKeyShared(t *testing.T) {
+	const MsgBatchCount = 100
+	client, err := NewClient(ClientOptions{
+		URL: lookupURL,
+	})
+	assert.Nil(t, err)
+	defer client.Close()
+
+	topic := "persistent://public/default/test-key-based-batch-with-key-shared"
+
+	consumer1, err := client.Subscribe(ConsumerOptions{
+		Topic:            topic,
+		SubscriptionName: "sub-1",
+		Type:             KeyShared,
+	})
+	assert.Nil(t, err)
+	defer consumer1.Close()
+
+	consumer2, err := client.Subscribe(ConsumerOptions{
+		Topic:            topic,
+		SubscriptionName: "sub-1",
+		Type:             KeyShared,
+	})
+	assert.Nil(t, err)
+	defer consumer2.Close()
+
+	// create producer
+	producer, err := client.CreateProducer(ProducerOptions{
+		Topic:               topic,
+		DisableBatching:     false,
+		BatcherBuilderType:  KeyBasedBatchBuilder,
+		BatchingMaxMessages: 10,
+	})
+	assert.Nil(t, err)
+	defer producer.Close()
+
+	ctx := context.Background()
+	keys := []string{"key1", "key2", "key3"}
+	for i := 0; i < MsgBatchCount; i++ {
+		for _, k := range keys {
+			producer.SendAsync(ctx, &ProducerMessage{
+				Key:     k,
+				Payload: []byte(fmt.Sprintf("value-%d", i)),
+			}, func(id MessageID, producerMessage *ProducerMessage, err error) {
+				assert.Nil(t, err)
+			},
+			)
+		}
+	}
+
+	receivedConsumer1 := 0
+	receivedConsumer2 := 0
+	consumer1Keys := make(map[string]int)
+	consumer2Keys := make(map[string]int)
+	for (receivedConsumer1 + receivedConsumer2) < 300 {
+		select {
+		case cm, ok := <-consumer1.Chan():
+			if !ok {
+				break
+			}
+			receivedConsumer1++
+			cnt := 0
+			if _, has := consumer1Keys[cm.Key()]; has {
+				cnt = consumer1Keys[cm.Key()]
+			}
+			assert.Equal(
+				t, fmt.Sprintf("value-%d", cnt),
+				string(cm.Payload()),
+			)
+			consumer1Keys[cm.Key()] = cnt + 1
+			consumer1.Ack(cm.Message)
+		case cm, ok := <-consumer2.Chan():
+			if !ok {
+				break
+			}
+			receivedConsumer2++
+			cnt := 0
+			if _, has := consumer2Keys[cm.Key()]; has {
+				cnt = consumer2Keys[cm.Key()]
+			}
+			assert.Equal(
+				t, fmt.Sprintf("value-%d", cnt),
+				string(cm.Payload()),
+			)
+			consumer2Keys[cm.Key()] = cnt + 1
+			consumer2.Ack(cm.Message)
+		}
+	}
+
+	assert.NotEqual(t, 0, receivedConsumer1)
+	assert.NotEqual(t, 0, receivedConsumer2)
+	assert.Equal(t, len(consumer1Keys)*MsgBatchCount, receivedConsumer1)
+	assert.Equal(t, len(consumer2Keys)*MsgBatchCount, receivedConsumer2)
+
+	fmt.Printf("TestKeyBasedBatchProducerConsumerKeyShared received messages consumer1: %d consumser2: %d\n",
+		receivedConsumer1, receivedConsumer2)
+	assert.Equal(t, 300, receivedConsumer1+receivedConsumer2)
+
+	fmt.Printf("TestKeyBasedBatchProducerConsumerKeyShared received messages keys consumer1: %v consumser2: %v\n",
+		consumer1Keys, consumer2Keys)
+}
+
+func TestOrderingOfKeyBasedBatchProducerConsumerKeyShared(t *testing.T) {
+	const MsgBatchCount = 10
+	client, err := NewClient(ClientOptions{
+		URL: lookupURL,
+	})
+	assert.Nil(t, err)
+	defer client.Close()
+
+	topic := "persistent://public/default/test-ordering-of-key-based-batch-with-key-shared"
+
+	consumer1, err := client.Subscribe(ConsumerOptions{
+		Topic:            topic,
+		SubscriptionName: "sub-1",
+		Type:             KeyShared,
+	})
+	assert.Nil(t, err)
+	defer consumer1.Close()
+
+	// create producer
+	producer, err := client.CreateProducer(ProducerOptions{
+		Topic:                   topic,
+		DisableBatching:         false,
+		BatcherBuilderType:      KeyBasedBatchBuilder,
+		BatchingMaxMessages:     30,
+		BatchingMaxPublishDelay: time.Second * 5,
+	})
+	assert.Nil(t, err)
+	defer producer.Close()
+
+	ctx := context.Background()
+	keys := []string{"key1", "key2", "key3"}
+	for i := 0; i < MsgBatchCount; i++ {
+		for _, k := range keys {
+			producer.SendAsync(ctx, &ProducerMessage{
+				Key:     k,
+				Payload: []byte(fmt.Sprintf("value-%d", i)),
+			}, func(id MessageID, producerMessage *ProducerMessage, err error) {
+				assert.Nil(t, err)
+			},
+			)
+		}
+	}
+
+	var receivedKey string
+	var receivedMessageIndex int
+	for i := 0; i < len(keys)*MsgBatchCount; i++ {
+		cm, ok := <-consumer1.Chan()
+		if !ok {
+			break
+		}
+		if receivedKey != cm.Key() {
+			receivedKey = cm.Key()
+			receivedMessageIndex = 0
+		}
+		assert.Equal(
+			t, fmt.Sprintf("value-%d", receivedMessageIndex%10),
+			string(cm.Payload()),
+		)
+		consumer1.Ack(cm.Message)
+		receivedMessageIndex++
+	}
+
+	// TODO: add OrderingKey support, see GH issue #401
+}