From 432fde88a92b8f9e95f546be503c08767cf15975 Mon Sep 17 00:00:00 2001
From: Wilson Wang <wilsonny371@gmail.com>
Date: Thu, 4 Mar 2021 11:02:00 -0800
Subject: [PATCH] pkg/wait: change list from single element to an array.

We found wait lock contention when a large amount of write operations. Converting wait from single element to an array helps to improve the performance.

Fixes #12731

Signed-off-by: Wilson Wang <wilsonny371@gmail.com>
---
 pkg/wait/wait.go | 49 +++++++++++++++++++++++++++++++++---------------
 1 file changed, 34 insertions(+), 15 deletions(-)

diff --git a/pkg/wait/wait.go b/pkg/wait/wait.go
index 9b1df419e..a3e2aec7c 100644
--- a/pkg/wait/wait.go
+++ b/pkg/wait/wait.go
@@ -21,6 +21,13 @@ import (
 	"sync"
 )
 
+const (
+	// To avoid lock contention we use an array of list struct (rw mutex & map)
+	// for the id argument, we apply mod operation and uses its remainder to
+	// index into the array and find the corresponding element.
+	defaultListElementLength = 64
+)
+
 // Wait is an interface that provides the ability to wait and trigger events that
 // are associated with IDs.
 type Wait interface {
@@ -34,33 +41,44 @@ type Wait interface {
 }
 
 type list struct {
+	e []listElement
+}
+
+type listElement struct {
 	l sync.RWMutex
 	m map[uint64]chan interface{}
 }
 
 // New creates a Wait.
 func New() Wait {
-	return &list{m: make(map[uint64]chan interface{})}
+	res := list{
+		e: make([]listElement, defaultListElementLength),
+	}
+	for i := 0; i < len(res.e); i++ {
+		res.e[i].m = make(map[uint64]chan interface{})
+	}
+	return &res
 }
 
 func (w *list) Register(id uint64) <-chan interface{} {
-	w.l.Lock()
-	defer w.l.Unlock()
-	ch := w.m[id]
-	if ch == nil {
-		ch = make(chan interface{}, 1)
-		w.m[id] = ch
+	idx := id % defaultListElementLength
+	newCh := make(chan interface{}, 1)
+	w.e[idx].l.Lock()
+	defer w.e[idx].l.Unlock()
+	if _, ok := w.e[idx].m[id]; !ok {
+		w.e[idx].m[id] = newCh
 	} else {
 		log.Panicf("dup id %x", id)
 	}
-	return ch
+	return newCh
 }
 
 func (w *list) Trigger(id uint64, x interface{}) {
-	w.l.Lock()
-	ch := w.m[id]
-	delete(w.m, id)
-	w.l.Unlock()
+	idx := id % defaultListElementLength
+	w.e[idx].l.Lock()
+	ch := w.e[idx].m[id]
+	delete(w.e[idx].m, id)
+	w.e[idx].l.Unlock()
 	if ch != nil {
 		ch <- x
 		close(ch)
@@ -68,9 +86,10 @@ func (w *list) Trigger(id uint64, x interface{}) {
 }
 
 func (w *list) IsRegistered(id uint64) bool {
-	w.l.RLock()
-	defer w.l.RUnlock()
-	_, ok := w.m[id]
+	idx := id % defaultListElementLength
+	w.e[idx].l.RLock()
+	defer w.e[idx].l.RUnlock()
+	_, ok := w.e[idx].m[id]
 	return ok
 }