nsqd: BenchmarkTopicMessagePump

Author: jehiah

nsqd/topic_test.go

@@ -1,13 +1,15 @@
 package nsqd
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"io/ioutil"
 	"net/http"
 	"os"
 	"runtime"
 	"strconv"
+	"sync"
 	"testing"
 	"time"
 
@@ -273,3 +275,42 @@ func BenchmarkTopicToChannelPut(b *testing.B) {
 		runtime.Gosched()
 	}
 }
+
+func BenchmarkTopicMessagePump(b *testing.B) {
+	b.StopTimer()
+	topicName := "bench_topic_put_throughput" + strconv.Itoa(b.N)
+	opts := NewOptions()
+	opts.Logger = test.NewTestLogger(b)
+	opts.LogLevel = LOG_WARN
+	opts.MemQueueSize = int64(b.N)
+	_, _, nsqd := mustStartNSQD(opts)
+	defer os.RemoveAll(opts.DataPath)
+	defer nsqd.Exit()
+
+	topic := nsqd.GetOrCreateTopic(topicName)
+	ch := topic.GetOrCreateChannel("ch")
+	ctx, cancel := context.WithCancel(context.Background())
+
+	var wg sync.WaitGroup
+	for i := 0; i < runtime.GOMAXPROCS(0); i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for {
+				select {
+				case <-ch.memoryMsgChan:
+				case <-ctx.Done():
+					return
+				}
+			}
+		}()
+	}
+
+	b.StartTimer()
+	for i := 0; i <= b.N; i++ {
+		msg := NewMessage(topic.GenerateID(), []byte("aaaaaaaaaaaaaaaaaaaaaaaaaaa"))
+		topic.PutMessage(msg)
+	}
+	cancel()
+	wg.Wait()
+}