Configurable parallelism for benchmarks (#333)

centrifugal · Nov 5, 2023 · 83312da · 83312da
1 parent 8ebee8a
commit 83312da
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 13 deletions.
diff --git a/broker_memory_test.go b/broker_memory_test.go
@@ -2,6 +2,7 @@ package centrifuge
 
 import (
 	"context"
+	"os"
 	"strconv"
 	"testing"
 	"time"
@@ -10,6 +11,20 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+const defaultParallelism = 128
+
+func getBenchParallelism() int {
+	parallelism := os.Getenv("PARALLELISM")
+	if parallelism == "" {
+		return defaultParallelism
+	}
+	p, err := strconv.Atoi(parallelism)
+	if err != nil {
+		panic(err)
+	}
+	return p
+}
+
 func testMemoryBroker() *MemoryBroker {
 	n, err := New(Config{
 		LogLevel:   LogLevelDebug,
@@ -294,7 +309,7 @@ func BenchmarkMemoryPublish_1Ch(b *testing.B) {
 	defer func() { _ = e.node.Shutdown(context.Background()) }()
 
 	rawData := protocol.Raw(`{"bench": true}`)
-	b.SetParallelism(128)
+	b.SetParallelism(getBenchParallelism())
 	b.ResetTimer()
 	b.RunParallel(func(pb *testing.PB) {
 		for pb.Next() {
@@ -312,7 +327,7 @@ func BenchmarkMemoryPublish_History_1Ch(b *testing.B) {
 
 	rawData := protocol.Raw(`{"bench": true}`)
 	chOpts := PublishOptions{HistorySize: 100, HistoryTTL: 60 * time.Second}
-	b.SetParallelism(128)
+	b.SetParallelism(getBenchParallelism())
 	b.ResetTimer()
 	b.RunParallel(func(pb *testing.PB) {
 		for pb.Next() {

diff --git a/broker_redis_test.go b/broker_redis_test.go
@@ -1222,7 +1222,7 @@ func BenchmarkRedisPublish_1Ch(b *testing.B) {
 			defer func() { _ = node.Shutdown(context.Background()) }()
 			defer stopRedisBroker(broker)
 			rawData := []byte(`{"bench": true}`)
-			b.SetParallelism(128)
+			b.SetParallelism(getBenchParallelism())
 			b.ResetTimer()
 			b.RunParallel(func(pb *testing.PB) {
 				for pb.Next() {
@@ -1246,7 +1246,7 @@ func BenchmarkRedisPublish_ManyCh(b *testing.B) {
 			defer func() { _ = node.Shutdown(context.Background()) }()
 			defer stopRedisBroker(broker)
 			rawData := []byte(`{"bench": true}`)
-			b.SetParallelism(128)
+			b.SetParallelism(getBenchParallelism())
 			j := int32(0)
 			b.ResetTimer()
 			b.RunParallel(func(pb *testing.PB) {
@@ -1272,7 +1272,7 @@ func BenchmarkRedisPublish_History_1Ch(b *testing.B) {
 			defer stopRedisBroker(broker)
 			rawData := []byte(`{"bench": true}`)
 			chOpts := PublishOptions{HistorySize: 100, HistoryTTL: 100 * time.Second}
-			b.SetParallelism(128)
+			b.SetParallelism(getBenchParallelism())
 			b.ResetTimer()
 			b.RunParallel(func(pb *testing.PB) {
 				for pb.Next() {
@@ -1299,7 +1299,7 @@ func BenchmarkRedisPub_History_ManyCh(b *testing.B) {
 			defer stopRedisBroker(broker)
 			rawData := []byte(`{"bench": true}`)
 			chOpts := PublishOptions{HistorySize: 100, HistoryTTL: 100 * time.Second}
-			b.SetParallelism(128)
+			b.SetParallelism(getBenchParallelism())
 			j := int32(0)
 			b.ResetTimer()
 			b.RunParallel(func(pb *testing.PB) {
@@ -1339,7 +1339,7 @@ func BenchmarkRedisSubscribe(b *testing.B) {
 			defer func() { _ = node.Shutdown(context.Background()) }()
 			defer stopRedisBroker(broker)
 			i := int32(0)
-			b.SetParallelism(128)
+			b.SetParallelism(getBenchParallelism())
 			b.ResetTimer()
 			b.RunParallel(func(pb *testing.PB) {
 				for pb.Next() {
@@ -1366,7 +1366,7 @@ func BenchmarkRedisHistory_1Ch(b *testing.B) {
 				_, err := broker.Publish("channel", rawData, PublishOptions{HistorySize: 4, HistoryTTL: 300 * time.Second})
 				require.NoError(b, err)
 			}
-			b.SetParallelism(128)
+			b.SetParallelism(getBenchParallelism())
 			b.ResetTimer()
 			b.RunParallel(func(pb *testing.PB) {
 				for pb.Next() {
@@ -1394,7 +1394,7 @@ func BenchmarkRedisRecover_1Ch(b *testing.B) {
 				_, err := broker.Publish("channel", rawData, PublishOptions{HistorySize: numMessages, HistoryTTL: 300 * time.Second})
 				require.NoError(b, err)
 			}
-			b.SetParallelism(128)
+			b.SetParallelism(getBenchParallelism())
 			b.ResetTimer()
 			b.RunParallel(func(pb *testing.PB) {
 				for pb.Next() {
@@ -1655,7 +1655,7 @@ func BenchmarkPubSubThroughput(b *testing.B) {
 			defer stopRedisBroker(b2)
 
 			b.ReportAllocs()
-			b.SetParallelism(128)
+			b.SetParallelism(getBenchParallelism())
 			b.ResetTimer()
 			var i int64
 			b.RunParallel(func(pb *testing.PB) {

diff --git a/presence_memory_test.go b/presence_memory_test.go
@@ -102,6 +102,7 @@ func BenchmarkMemoryAddPresence_OneChannel_Parallel(b *testing.B) {
 	defer func() { _ = e.node.Shutdown(context.Background()) }()
 
 	b.ResetTimer()
+	b.SetParallelism(getBenchParallelism())
 	b.RunParallel(func(pb *testing.PB) {
 		for pb.Next() {
 			err := e.AddPresence("channel", "uid", &ClientInfo{})
@@ -131,6 +132,7 @@ func BenchmarkMemoryPresence_OneChannel_Parallel(b *testing.B) {
 	defer func() { _ = e.node.Shutdown(context.Background()) }()
 
 	_ = e.AddPresence("channel", "uid", &ClientInfo{})
+	b.SetParallelism(getBenchParallelism())
 	b.ResetTimer()
 	b.RunParallel(func(pb *testing.PB) {
 		for pb.Next() {

diff --git a/presence_redis_test.go b/presence_redis_test.go
@@ -111,7 +111,7 @@ func BenchmarkRedisAddPresence_1Ch(b *testing.B) {
 			pm := newTestRedisPresenceManager(b, node, tt.UseCluster)
 			defer func() { _ = node.Shutdown(context.Background()) }()
 			defer stopRedisPresenceManager(pm)
-			b.SetParallelism(128)
+			b.SetParallelism(getBenchParallelism())
 			b.ResetTimer()
 			b.RunParallel(func(pb *testing.PB) {
 				for pb.Next() {
@@ -132,7 +132,7 @@ func BenchmarkRedisPresence_1Ch(b *testing.B) {
 			pm := newTestRedisPresenceManager(b, node, tt.UseCluster)
 			defer func() { _ = node.Shutdown(context.Background()) }()
 			defer stopRedisPresenceManager(pm)
-			b.SetParallelism(128)
+			b.SetParallelism(getBenchParallelism())
 			_ = pm.AddPresence("channel", "uid", &ClientInfo{})
 			b.ResetTimer()
 			b.RunParallel(func(pb *testing.PB) {
@@ -154,7 +154,7 @@ func BenchmarkRedisPresence_ManyCh(b *testing.B) {
 			pm := newTestRedisPresenceManager(b, node, tt.UseCluster)
 			defer func() { _ = node.Shutdown(context.Background()) }()
 			defer stopRedisPresenceManager(pm)
-			b.SetParallelism(128)
+			b.SetParallelism(getBenchParallelism())
 			_ = pm.AddPresence("channel", "uid", &ClientInfo{})
 			j := int32(0)
 			b.ResetTimer()