perf: reduce contention caused by the usage of runtime.GOMAXPROCS

Author: rueian

internal/util/parallel.go

@@ -1,41 +1,40 @@
 package util
 
 import (
-	"runtime"
 	"sync"
 )
 
-func ParallelKeys[K comparable, V any](p map[K]V, fn func(k K)) {
+func ParallelKeys[K comparable, V any](maxp int, p map[K]V, fn func(k K)) {
 	ch := make(chan K, len(p))
 	for k := range p {
 		ch <- k
 	}
-	closeThenParallel(ch, fn)
+	closeThenParallel(maxp, ch, fn)
 }
 
-func ParallelVals[K comparable, V any](p map[K]V, fn func(k V)) {
+func ParallelVals[K comparable, V any](maxp int, p map[K]V, fn func(k V)) {
 	ch := make(chan V, len(p))
 	for _, v := range p {
 		ch <- v
 	}
-	closeThenParallel(ch, fn)
+	closeThenParallel(maxp, ch, fn)
 }
 
-func closeThenParallel[V any](ch chan V, fn func(k V)) {
+func closeThenParallel[V any](maxp int, ch chan V, fn func(k V)) {
 	close(ch)
 	concurrency := len(ch)
-	if cpus := runtime.GOMAXPROCS(0); concurrency > cpus {
-		concurrency = cpus
+	if concurrency > maxp {
+		concurrency = maxp
 	}
 	wg := sync.WaitGroup{}
 	wg.Add(concurrency)
 	for i := 1; i < concurrency; i++ {
-		go func() {
+		go func(wg *sync.WaitGroup) {
 			for v := range ch {
 				fn(v)
 			}
 			wg.Done()
-		}()
+		}(&wg)
 	}
 	for v := range ch {
 		fn(v)

internal/util/parallel_test.go

@@ -9,7 +9,7 @@ import (
 func TestParallelKeys(t *testing.T) {
 	var sum int64
 	data, sk, _ := gen(int64(runtime.NumCPU() * 1000))
-	ParallelKeys(data, func(i int64) {
+	ParallelKeys(runtime.GOMAXPROCS(0), data, func(i int64) {
 		atomic.AddInt64(&sum, i)
 	})
 	if atomic.LoadInt64(&sum) != sk {
@@ -20,7 +20,7 @@ func TestParallelKeys(t *testing.T) {
 func TestParallelVals(t *testing.T) {
 	var sum int64
 	data, _, sv := gen(int64(runtime.NumCPU() * 1000))
-	ParallelVals(data, func(i int64) {
+	ParallelVals(runtime.GOMAXPROCS(0), data, func(i int64) {
 		atomic.AddInt64(&sum, i)
 	})
 	if atomic.LoadInt64(&sum) != sv {

lua.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"crypto/sha1"
 	"encoding/hex"
+	"runtime"
 	"sync/atomic"
 
 	"github.com/redis/rueidis/internal/util"
@@ -12,7 +13,7 @@ import (
 // NewLuaScript creates a Lua instance whose Lua.Exec uses EVALSHA and EVAL.
 func NewLuaScript(script string) *Lua {
 	sum := sha1.Sum([]byte(script))
-	return &Lua{script: script, sha1: hex.EncodeToString(sum[:])}
+	return &Lua{script: script, sha1: hex.EncodeToString(sum[:]), maxp: runtime.GOMAXPROCS(0)}
 }
 
 // NewLuaScriptReadOnly creates a Lua instance whose Lua.Exec uses EVALSHA_RO and EVAL_RO.
@@ -26,6 +27,7 @@ func NewLuaScriptReadOnly(script string) *Lua {
 type Lua struct {
 	script   string
 	sha1     string
+	maxp     int
 	readonly bool
 }
 
@@ -59,7 +61,7 @@ type LuaExec struct {
 // Cross slot keys within single LuaExec are prohibited if the Client is a cluster client.
 func (s *Lua) ExecMulti(ctx context.Context, c Client, multi ...LuaExec) (resp []RedisResult) {
 	var e atomic.Value
-	util.ParallelVals(c.Nodes(), func(n Client) {
+	util.ParallelVals(s.maxp, c.Nodes(), func(n Client) {
 		if err := n.Do(ctx, n.B().ScriptLoad().Script(s.script).Build()).Error(); err != nil {
 			e.CompareAndSwap(nil, &errs{error: err})
 		}

mux.go

@@ -3,6 +3,7 @@ package rueidis
 import (
 	"context"
 	"net"
+	"runtime"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -48,6 +49,7 @@ type mux struct {
 	wire   []atomic.Value
 	sc     []*singleconnect
 	mu     []sync.Mutex
+	maxp   int
 }
 
 func makeMux(dst string, option *ClientOption, dialFn dialFn) *mux {
@@ -75,6 +77,7 @@ func newMux(dst string, option *ClientOption, init, dead wire, wireFn wireFn) *m
 		wire: make([]atomic.Value, multiplex),
 		mu:   make([]sync.Mutex, multiplex),
 		sc:   make([]*singleconnect, multiplex),
+		maxp: runtime.GOMAXPROCS(0),
 	}
 	for i := 0; i < len(m.wire); i++ {
 		m.wire[i].Store(init)
@@ -260,7 +263,7 @@ func (m *mux) DoMultiCache(ctx context.Context, multi ...CacheableTTL) (results
 	}
 
 	results = make([]RedisResult, len(multi))
-	util.ParallelKeys(commands, func(slot uint16) {
+	util.ParallelKeys(m.maxp, commands, func(slot uint16) {
 		for i, r := range m.doMultiCache(ctx, slot, commands[slot]) {
 			results[cIndexes[slot][i]] = r
 		}

rueidiscompat/adapter.go

@@ -30,6 +30,7 @@ import (
 	"context"
 	"encoding"
 	"fmt"
+	"runtime"
 	"strconv"
 	"strings"
 	"sync"
@@ -389,6 +390,7 @@ type Cmdable interface {
 
 type Compat struct {
 	client rueidis.Client
+	maxp   int
 }
 
 type CacheCompat struct {
@@ -397,7 +399,7 @@ type CacheCompat struct {
 }
 
 func NewAdapter(client rueidis.Client) Cmdable {
-	return &Compat{client: client}
+	return &Compat{client: client, maxp: runtime.GOMAXPROCS(0)}
 }
 
 func (c *Compat) Cache(ttl time.Duration) CacheCompat {
@@ -2799,7 +2801,7 @@ func (c *Compat) ACLDryRun(ctx context.Context, username string, command ...any)
 
 func (c *Compat) doPrimaries(ctx context.Context, fn func(c rueidis.Client) error) error {
 	var firsterr atomic.Value
-	util.ParallelVals(c.client.Nodes(), func(client rueidis.Client) {
+	util.ParallelVals(c.maxp, c.client.Nodes(), func(client rueidis.Client) {
 		msgs, err := client.Do(ctx, client.B().Role().Build()).ToArray()
 		if err == nil {
 			if role, _ := msgs[0].ToString(); role == "master" {