From 75d24d4c36b11663ae88601e4713cf76a4eb059b Mon Sep 17 00:00:00 2001
From: Sam Xie <xsambundy@gmail.com>
Date: Thu, 16 Jul 2020 16:28:14 +0800
Subject: [PATCH] Add benchmarks for sarama instrumentation

---
 .../Shopify/sarama/consumer_test.go           | 43 ++++++++++
 .../Shopify/sarama/producer_test.go           | 82 +++++++++++++++++++
 2 files changed, 125 insertions(+)

diff --git a/instrumentation/github.com/Shopify/sarama/consumer_test.go b/instrumentation/github.com/Shopify/sarama/consumer_test.go
index 8b1658ab3a5..414ed5dc851 100644
--- a/instrumentation/github.com/Shopify/sarama/consumer_test.go
+++ b/instrumentation/github.com/Shopify/sarama/consumer_test.go
@@ -165,3 +165,46 @@ func TestConsumer_ConsumePartitionWithError(t *testing.T) {
 	_, err = consumer.ConsumePartition(topic, 0, 0)
 	assert.Error(t, err)
 }
+
+func BenchmarkWrapPartitionConsumer(b *testing.B) {
+	// Mock tracer
+	mt := mocktracer.NewTracer("kafka")
+
+	mockPartitionConsumer, partitionConsumer := createMockPartitionConsumer(b)
+
+	partitionConsumer = WrapPartitionConsumer(serviceName, partitionConsumer, WithTracer(mt))
+	message := sarama.ConsumerMessage{Key: []byte("foo")}
+
+	b.ReportAllocs()
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		mockPartitionConsumer.YieldMessage(&message)
+		<-partitionConsumer.Messages()
+	}
+}
+
+func BenchmarkMockPartitionConsumer(b *testing.B) {
+	mockPartitionConsumer, partitionConsumer := createMockPartitionConsumer(b)
+
+	message := sarama.ConsumerMessage{Key: []byte("foo")}
+
+	b.ReportAllocs()
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		mockPartitionConsumer.YieldMessage(&message)
+		<-partitionConsumer.Messages()
+	}
+}
+
+func createMockPartitionConsumer(b *testing.B) (*mocks.PartitionConsumer, sarama.PartitionConsumer) {
+	// Mock partition consumer controller
+	consumer := mocks.NewConsumer(b, sarama.NewConfig())
+	mockPartitionConsumer := consumer.ExpectConsumePartition(topic, 0, 0)
+
+	// Create partition consumer
+	partitionConsumer, err := consumer.ConsumePartition(topic, 0, 0)
+	assert.NoError(b, err)
+	return mockPartitionConsumer, partitionConsumer
+}
diff --git a/instrumentation/github.com/Shopify/sarama/producer_test.go b/instrumentation/github.com/Shopify/sarama/producer_test.go
index c2bc2400563..25707f13c29 100644
--- a/instrumentation/github.com/Shopify/sarama/producer_test.go
+++ b/instrumentation/github.com/Shopify/sarama/producer_test.go
@@ -328,3 +328,85 @@ func newSaramaConfig() *sarama.Config {
 	cfg.Version = sarama.V0_11_0_0
 	return cfg
 }
+
+func BenchmarkWrapSyncProducer(b *testing.B) {
+	// Mock tracer
+	mt := mocktracer.NewTracer("kafka")
+
+	cfg := newSaramaConfig()
+	// Mock sync producer
+	mockSyncProducer := mocks.NewSyncProducer(b, cfg)
+
+	// Wrap sync producer
+	syncProducer := WrapSyncProducer(serviceName, cfg, mockSyncProducer, WithTracer(mt))
+	message := sarama.ProducerMessage{Key: sarama.StringEncoder("foo")}
+
+	b.ReportAllocs()
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		mockSyncProducer.ExpectSendMessageAndSucceed()
+		_, _, err := syncProducer.SendMessage(&message)
+		assert.NoError(b, err)
+	}
+}
+
+func BenchmarkMockSyncProducer(b *testing.B) {
+	cfg := newSaramaConfig()
+	// Mock sync producer
+	mockSyncProducer := mocks.NewSyncProducer(b, cfg)
+
+	// Wrap sync producer
+	syncProducer := mockSyncProducer
+	message := sarama.ProducerMessage{Key: sarama.StringEncoder("foo")}
+
+	b.ReportAllocs()
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		mockSyncProducer.ExpectSendMessageAndSucceed()
+		_, _, err := syncProducer.SendMessage(&message)
+		assert.NoError(b, err)
+	}
+}
+
+func BenchmarkWrapAsyncProducer(b *testing.B) {
+	// Mock tracer
+	mt := mocktracer.NewTracer("kafka")
+
+	cfg := newSaramaConfig()
+	cfg.Producer.Return.Successes = true
+	mockAsyncProducer := mocks.NewAsyncProducer(b, cfg)
+
+	// Wrap sync producer
+	asyncProducer := WrapAsyncProducer(serviceName, cfg, mockAsyncProducer, WithTracer(mt))
+	message := sarama.ProducerMessage{Key: sarama.StringEncoder("foo")}
+
+	b.ReportAllocs()
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		mockAsyncProducer.ExpectInputAndSucceed()
+		asyncProducer.Input() <- &message
+		<-asyncProducer.Successes()
+	}
+}
+
+func BenchmarkMockAsyncProducer(b *testing.B) {
+	cfg := newSaramaConfig()
+	cfg.Producer.Return.Successes = true
+	mockAsyncProducer := mocks.NewAsyncProducer(b, cfg)
+
+	// Wrap sync producer
+	asyncProducer := mockAsyncProducer
+	message := sarama.ProducerMessage{Key: sarama.StringEncoder("foo")}
+
+	b.ReportAllocs()
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		mockAsyncProducer.ExpectInputAndSucceed()
+		mockAsyncProducer.Input() <- &message
+		<-asyncProducer.Successes()
+	}
+}