Add single partition router

pulsar/producer_test.go

@@ -560,6 +560,54 @@ func TestMessageRouter(t *testing.T) {
 	assert.NotNil(t, msg)
 	assert.Equal(t, string(msg.Payload()), "hello")
 }
+func TestMessageSingleRouter(t *testing.T) {
+	// Create topic with 5 partitions
+	topicAdminURL := "admin/v2/persistent/public/default/my-single-partitioned-topic/partitions"
+	err := httpPut(topicAdminURL, 5)
+	defer httpDelete(topicAdminURL)
+	if err != nil {
+		t.Fatal(err)
+	}
+	client, err := NewClient(ClientOptions{
+		URL: serviceURL,
+	})
+
+	assert.Nil(t, err)
+	defer client.Close()
+
+	// Only subscribe on the specific partition
+	consumer, err := client.Subscribe(ConsumerOptions{
+		Topic:            "my-single-partitioned-topic",
+		SubscriptionName: "my-sub",
+	})
+
+	assert.Nil(t, err)
+	defer consumer.Close()
+
+	producer, err := client.CreateProducer(ProducerOptions{
+		Topic:         "my-single-partitioned-topic",
+		MessageRouter: NewSinglePartitionRouter(),
+	})
+
+	assert.Nil(t, err)
+	defer producer.Close()
+
+	ctx := context.Background()
+
+	ID, err := producer.Send(ctx, &ProducerMessage{
+		Payload: []byte("hello"),
+	})
+	assert.Nil(t, err)
+	assert.NotNil(t, ID)
+
+	fmt.Println("PUBLISHED")
+
+	// Verify message was published on partition 2
+	msg, err := consumer.Receive(ctx)
+	assert.Nil(t, err)
+	assert.NotNil(t, msg)
+	assert.Equal(t, string(msg.Payload()), "hello")
+}
 
 func TestNonPersistentTopic(t *testing.T) {
 	topicName := "non-persistent://public/default/testNonPersistentTopic"

pulsar/single_partition_route_bench_test.go

@@ -0,0 +1,39 @@
+// 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 (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func BenchmarkSinglePartitionRouter(b *testing.B) {
+	numPartitions := topicMetaData{300}
+	router := NewSinglePartitionRouter()
+	var expect *int
+	for i := 0; i < b.N; i++ {
+		p := router(&ProducerMessage{
+			Payload: []byte("message 2"),
+		}, numPartitions)
+		if expect == nil {
+			expect = &p
+		}
+		assert.Equal(b, *expect, p)
+	}
+}

pulsar/single_partition_router.go

@@ -0,0 +1,50 @@
+// 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 "sync"
+
+var (
+	singlePartition *int
+	mutex           sync.Mutex
+)
+
+func NewSinglePartitionRouter() func(*ProducerMessage, TopicMetadata) int {
+	return func(message *ProducerMessage, metadata TopicMetadata) int {
+		numPartitions := metadata.NumPartitions()
+		if len(message.OrderingKey) != 0 {
+			// When an OrderingKey is specified, use the hash of that key
+			return int(getHashingFunction(JavaStringHash)(message.OrderingKey) % numPartitions)
+		}
+
+		if len(message.Key) != 0 {
+			// When a key is specified, use the hash of that key
+			return int(getHashingFunction(JavaStringHash)(message.Key) % numPartitions)
+		}
+		if singlePartition != nil {
+			return *singlePartition
+		}
+		mutex.Lock()
+		defer mutex.Unlock()
+		partition := r.R.Intn(int(numPartitions))
+		singlePartition = &partition
+		return partition
+
+	}
+
+}

pulsar/single_partition_router_test.go

@@ -0,0 +1,77 @@
+// 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 (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+type topicMetaData struct {
+	partition uint32
+}
+
+func (t topicMetaData) NumPartitions() uint32 {
+	return t.partition
+}
+
+func TestNewSinglePartitionRouter(t *testing.T) {
+	numPartitions := topicMetaData{2}
+	router := NewSinglePartitionRouter()
+	p := router(&ProducerMessage{
+		Payload: []byte("message 2"),
+	}, numPartitions)
+	assert.GreaterOrEqual(t, p, 0)
+
+	p2 := router(&ProducerMessage{
+		Payload: []byte("message 2"),
+	}, numPartitions)
+	assert.Equal(t, p, p2)
+}
+
+func TestNewSinglePartitionRouterWithKey(t *testing.T) {
+	router := NewSinglePartitionRouter()
+	numPartitions := topicMetaData{3}
+	p := router(&ProducerMessage{
+		Payload: []byte("message 2"),
+		Key:     "my-key",
+	}, numPartitions)
+	assert.Equal(t, 1, p)
+
+	p2 := router(&ProducerMessage{
+		Key:     "my-key",
+		Payload: []byte("message 2"),
+	}, numPartitions)
+	assert.Equal(t, p, p2)
+}
+func TestNewSinglePartitionRouterWithOrderingKey(t *testing.T) {
+	router := NewSinglePartitionRouter()
+	numPartitions := topicMetaData{3}
+	p := router(&ProducerMessage{
+		Payload:     []byte("message 2"),
+		OrderingKey: "my-key",
+	}, numPartitions)
+	assert.Equal(t, 1, p)
+
+	p2 := router(&ProducerMessage{
+		OrderingKey: "my-key",
+		Payload:     []byte("message 2"),
+	}, numPartitions)
+	assert.Equal(t, p, p2)
+}