gossip: parallelize callback handling

This commit introduces parallel callback handling in the gossip package
by:

1. Moving callback work management from a single global queue to
   per-callback-registration work queues.
2. Launching a dedicated worker goroutine for each callback
   registration.
3. Maintaining the guarantee that callbacks for the same key are
   executed sequentially. This only applies per callback registration,
   and there are no ordering guarantees between different callback
   registrations.

The changes improve performance by allowing different callbacks to
execute concurrently while preserving the ordering guarantees. Each
registration now has its own worker goroutine, channel for signaling new
work, and mutex-protected work queue.

Also, it changes the way the storeGossiper used to wait for gossip to
propagate. The storeGossiper used to subscribe to store updates, and
then it used to rely on the single threaded execution of gossip
callbacks to wait for its callback to get called, which means that all
other callbacks have been called since the storeGossiper subscribed
later. However, with parallel callback execution across different
subscriptions, we needed to introduce a more robust way to ensure the
propagation of the gossiped info. We do this by introducing a new
gossip method that creates a WaitGroup, incrementing it for every
callback match. the WaitGroup will get decremented at the end of
callbacks execution, ensuring that our update has definitely been
propagated.

Benchmark results before/after:

```
name                    old time/op    new time/op    delta
CallbackParallelism-12    5.69ms ± 0%    1.15ms ± 0%  -79.87%  (p=0.000 n=14+12)

name                    old alloc/op   new alloc/op   delta
CallbackParallelism-12      696B ±68%      897B ±12%  +28.83%  (p=0.020 n=15+15)

name                    old allocs/op  new allocs/op  delta
CallbackParallelism-12      10.0 ± 0%      13.0 ± 0%  +30.00%  (p=0.000 n=15+12)
```

Fixes: #144187

Release note: None

Author: iskettaneh

pkg/gossip/gossip.go

@@ -1565,6 +1565,66 @@ func (g *Gossip) findClient(match func(*client) bool) *client {
 	return nil
 }
 
+// TestingAddInfoProtoAndWaitForAllCallbacks adds an info proto, and waits for all
+// matching callbacks to get called before returning. It's only intended to be
+// used for tests that assert on the result of the gossip propagation.
+func (g *Gossip) TestingAddInfoProtoAndWaitForAllCallbacks(
+	key string, msg protoutil.Message, ttl time.Duration,
+) error {
+	// Take the lock to avoid races where a callback could be added while this
+	// method is waiting for matching callbacks to be called.
+	g.mu.Lock()
+	defer g.mu.Unlock()
+
+	wg := &sync.WaitGroup{}
+
+	// Increment the wait group once per matching callback. It will be decremented
+	// once the processing is complete.
+	for _, cb := range g.mu.is.callbacks {
+		if cb.matcher.MatchString(key) {
+			wg.Add(1)
+		}
+	}
+
+	// Add the target info to the infoStore. This will trigger the registered
+	// callbacks to be called.
+	bytes, err := protoutil.Marshal(msg)
+	if err != nil {
+		return err
+	}
+	if err := g.addInfoLocked(key, bytes, ttl); err != nil {
+		return err
+	}
+
+	// At this point, we know that the callbacks that will be called have been
+	// added to the work queues. Now, we can append an entry item at the end of
+	// the matching callback's work queue that will decrement the wait group that
+	// was incremented earlier. This ensures that ALL matching callbacks have
+	// been called.
+	for _, cb := range g.mu.is.callbacks {
+		if cb.matcher.MatchString(key) {
+			cb.cw.mu.Lock()
+			cb.cw.mu.workQueue = append(cb.cw.mu.workQueue, callbackWorkItem{
+				method: func(_ string, _ roachpb.Value) {
+					wg.Done()
+				},
+				schedulingTime: timeutil.Now(),
+			})
+			cb.cw.mu.Unlock()
+		}
+
+		// Make sure to notify the callback worker that there is work to do.
+		select {
+		case cb.cw.callbackCh <- struct{}{}:
+		default:
+		}
+	}
+
+	// Wait for all the callbacks to finish processing.
+	wg.Wait()
+	return nil
+}
+
 // A firstRangeMissingError indicates that the first range has not yet
 // been gossiped. This will be the case for a node which hasn't yet
 // joined the gossip network.

pkg/gossip/infostore.go

@@ -36,11 +36,76 @@ func (allMatcher) MatchString(string) bool {
 	return true
 }
 
-// callback holds regexp pattern match and GossipCallback method.
+// This is how multiple callback registrations are handled:
+/*
++------------------+     +------------------+     +------------------+
+| Callback Reg 1   |     | Callback Reg 2   | ... | Callback Reg N   |
+| Pattern: "key1.*"|     | Pattern: "key2.*"| ... | Pattern: "keyN.*"|
++------------------+     +------------------+     +------------------+
+        |                        |                        |
+        v                        v                        v
++------------------+     +------------------+     +------------------+
+| callbackWorker 1 |     | callbackWorker 2 | ... | callbackWorker N |
+| - callbackCh     |     | - callbackCh     | ... | - callbackCh     |
+| - stopperCh      |     | - stopperCh      | ... | - stopperCh      |
+| - workQueue      |     | - workQueue      | ... | - workQueue      |
++------------------+     +------------------+     +------------------+
+        |                        |                        |
+        v                        v                        v
++------------------+     +------------------+     +------------------+
+| Worker Goroutine |     | Worker Goroutine | ... | Worker Goroutine |
+| 1                |     | 2                | ... | N                |
++------------------+     +------------------+     +------------------+
+*/
+// When a new info is added, it is checked against all the callback matchers.
+// The work is added inside each workQueue which matches the info key.
+// Each worker goroutine independently processes its own workQueue in FIFO
+// order.
+
+// callback holds regexp pattern match, GossipCallback method, and a queue of
+// remaining work items.
 type callback struct {
 	matcher   stringMatcher
 	method    Callback
 	redundant bool
+	// cw contains all the information needed to orchestrate, schedule, and run
+	// callbacks for this specific matcher.
+	cw *callbackWork
+}
+
+// callbackWorkItem is a struct that contains the information needed to run
+// a callback.
+type callbackWorkItem struct {
+	// schedulingTime is the time when the callback was scheduled.
+	schedulingTime time.Time
+	method         Callback
+	key            string
+	content        roachpb.Value
+}
+
+type callbackWork struct {
+	// callbackCh is used to signal the callback worker to execute the work.
+	callbackCh chan struct{}
+	// stopperCh is used to signal the callback worker to stop.
+	stopperCh chan struct{}
+	mu        struct {
+		syncutil.Mutex
+		// workQueue contains the queue of callbacks that need to be called.
+		workQueue []callbackWorkItem
+	}
+}
+
+func newCallbackWork() *callbackWork {
+	return &callbackWork{
+		callbackCh: make(chan struct{}, 1),
+		stopperCh:  make(chan struct{}, 1),
+		mu: struct {
+			syncutil.Mutex
+			workQueue []callbackWorkItem
+		}{
+			workQueue: make([]callbackWorkItem, 0),
+		},
+	}
 }
 
 // infoStore objects manage maps of Info objects. They maintain a
@@ -64,10 +129,6 @@ type infoStore struct {
 	NodeAddr        util.UnresolvedAddr      `json:"-"`               // Address of node owning this info store: "host:port"
 	highWaterStamps map[roachpb.NodeID]int64 // Per-node information for gossip peers
 	callbacks       []*callback
-
-	callbackWorkMu syncutil.Mutex // Protects callbackWork
-	callbackWork   []func()
-	callbackCh     chan struct{} // Channel to signal the callback goroutine
 }
 
 var monoTime struct {
@@ -167,34 +228,58 @@ func newInfoStore(
 		Infos:           make(infoMap),
 		NodeAddr:        nodeAddr,
 		highWaterStamps: map[roachpb.NodeID]int64{},
-		callbackCh:      make(chan struct{}, 1),
 	}
+	return is
+}
 
+// launchCallbackWorker launches a worker goroutine that is responsible for
+// executing callbacks for one registered callback pattern.
+func (is *infoStore) launchCallbackWorker(ambient log.AmbientContext, cw *callbackWork) {
 	bgCtx := ambient.AnnotateCtx(context.Background())
-	_ = is.stopper.RunAsyncTask(bgCtx, "infostore", func(ctx context.Context) {
+	_ = is.stopper.RunAsyncTask(bgCtx, "callback worker", func(ctx context.Context) {
 		for {
 			for {
-				is.callbackWorkMu.Lock()
-				work := is.callbackWork
-				is.callbackWork = nil
-				is.callbackWorkMu.Unlock()
+				cw.mu.Lock()
+				wq := cw.mu.workQueue
+				cw.mu.workQueue = nil
+				cw.mu.Unlock()
 
-				if len(work) == 0 {
+				if len(wq) == 0 {
 					break
 				}
-				for _, w := range work {
-					w()
+
+				// Execute all the callbacks in the queue, making sure to update the
+				// metrics accordingly.
+				for _, work := range wq {
+					afterQueue := timeutil.Now()
+					queueDur := afterQueue.Sub(work.schedulingTime)
+					is.metrics.CallbacksPending.Dec(1)
+					if queueDur >= minCallbackDurationToRecord {
+						is.metrics.CallbacksPendingDuration.RecordValue(queueDur.Nanoseconds())
+					}
+
+					work.method(work.key, work.content)
+
+					afterProcess := timeutil.Now()
+					processDur := afterProcess.Sub(afterQueue)
+					is.metrics.CallbacksProcessed.Inc(1)
+					if processDur > minCallbackDurationToRecord {
+						is.metrics.CallbacksProcessingDuration.RecordValue(processDur.Nanoseconds())
+					}
+					afterQueue = afterProcess // update for next iteration
 				}
 			}
 
 			select {
-			case <-is.callbackCh:
+			case <-cw.callbackCh:
+				// New work has just arrived.
 			case <-is.stopper.ShouldQuiesce():
 				return
+			case <-cw.stopperCh:
+				return
 			}
 		}
 	})
-	return is
 }
 
 // newInfo allocates and returns a new info object using the specified
@@ -312,17 +397,24 @@ func (is *infoStore) registerCallback(
 		opt.apply(cb)
 	}
 
+	cb.cw = newCallbackWork()
+	is.launchCallbackWorker(is.AmbientContext, cb.cw)
 	is.callbacks = append(is.callbacks, cb)
+
 	if err := is.visitInfos(func(key string, i *Info) error {
 		if matcher.MatchString(key) {
-			is.runCallbacks(key, i.Value, method)
+			is.runCallbacks(key, i.Value, cb)
 		}
 		return nil
 	}, true /* deleteExpired */); err != nil {
 		panic(err)
 	}
 
 	return func() {
+		// Stop the callback worker's goroutine.
+		cb.cw.stopperCh <- struct{}{}
+
+		// Remove the callback from the list.
 		for i, targetCB := range is.callbacks {
 			if targetCB == cb {
 				numCBs := len(is.callbacks)
@@ -339,51 +431,41 @@ func (is *infoStore) registerCallback(
 // matching each callback's regular expression against the key and invoking
 // the corresponding callback method on a match.
 func (is *infoStore) processCallbacks(key string, content roachpb.Value, changed bool) {
-	var matches []Callback
+	var callbacks []*callback
 	for _, cb := range is.callbacks {
 		if (changed || cb.redundant) && cb.matcher.MatchString(key) {
-			matches = append(matches, cb.method)
+			callbacks = append(callbacks, cb)
 		}
 	}
-	is.runCallbacks(key, content, matches...)
+	is.runCallbacks(key, content, callbacks...)
 }
 
-func (is *infoStore) runCallbacks(key string, content roachpb.Value, callbacks ...Callback) {
+// runCallbacks receives a list of callbacks and contents that match the key.
+// It adds work to the callback work slices, and signals the associated callback
+// workers to execute the work.
+func (is *infoStore) runCallbacks(key string, content roachpb.Value, callbacks ...*callback) {
 	// Add the callbacks to the callback work list.
 	beforeQueue := timeutil.Now()
 	is.metrics.CallbacksPending.Inc(int64(len(callbacks)))
-	f := func() {
-		afterQueue := timeutil.Now()
-		for _, method := range callbacks {
-			queueDur := afterQueue.Sub(beforeQueue)
-			is.metrics.CallbacksPending.Dec(1)
-			if queueDur >= minCallbackDurationToRecord {
-				is.metrics.CallbacksPendingDuration.RecordValue(queueDur.Nanoseconds())
-			}
-
-			method(key, content)
-
-			afterProcess := timeutil.Now()
-			processDur := afterProcess.Sub(afterQueue)
-			is.metrics.CallbacksProcessed.Inc(1)
-			if processDur > minCallbackDurationToRecord {
-				is.metrics.CallbacksProcessingDuration.RecordValue(processDur.Nanoseconds())
-			}
-			afterQueue = afterProcess // update for next iteration
+	for _, cb := range callbacks {
+		cb.cw.mu.Lock()
+		cb.cw.mu.workQueue = append(cb.cw.mu.workQueue, callbackWorkItem{
+			schedulingTime: beforeQueue,
+			method:         cb.method,
+			key:            key,
+			content:        content,
+		})
+		cb.cw.mu.Unlock()
+
+		// Signal the associated callback worker goroutine. Callbacks run in a
+		// goroutine to avoid mutex reentry. We also guarantee callbacks are run
+		// in order such that if a key is updated twice in succession, the second
+		// callback will never be run before the first.
+		select {
+		case cb.cw.callbackCh <- struct{}{}:
+		default:
 		}
 	}
-	is.callbackWorkMu.Lock()
-	is.callbackWork = append(is.callbackWork, f)
-	is.callbackWorkMu.Unlock()
-
-	// Signal the callback goroutine. Callbacks run in a goroutine to avoid mutex
-	// reentry. We also guarantee callbacks are run in order such that if a key
-	// is updated twice in succession, the second callback will never be run
-	// before the first.
-	select {
-	case is.callbackCh <- struct{}{}:
-	default:
-	}
 }
 
 // visitInfos implements a visitor pattern to run the visitInfo function against

pkg/gossip/infostore_test.go

@@ -688,8 +688,6 @@ func TestCallbacksCalledSequentially(t *testing.T) {
 // BenchmarkCallbackParallelism benchmarks the parallelism of the callback
 // worker. It registers multiple callbacks, and executes a fake workload
 // that sleeps for a short duration to simulate work done in the callback.
-// If we implement a parallel execution of the callback workers, we should
-// see a significant speedup.
 func BenchmarkCallbackParallelism(b *testing.B) {
 	ctx := context.Background()
 	is, stopper := newTestInfoStore()

pkg/kv/kvserver/allocator/allocatorimpl/allocator_test.go

@@ -9272,10 +9272,6 @@ func exampleRebalancing(
 	}, nil)
 
 	var wg sync.WaitGroup
-	g.RegisterCallback(gossip.MakePrefixPattern(gossip.KeyStoreDescPrefix),
-		func(_ string, _ roachpb.Value) { wg.Done() },
-		// Redundant callbacks are required by this test.
-		gossip.Redundant)
 
 	// Initialize testStores.
 	initTestStores(
@@ -9298,13 +9294,12 @@ func exampleRebalancing(
 	const generations = 100
 	for i := 0; i < generations; i++ {
 		// First loop through test stores and add data.
-		wg.Add(len(testStores))
 		for j := 0; j < len(testStores); j++ {
 			// Add a pretend range to the testStore if there's already one.
 			if testStores[j].Capacity.RangeCount > 0 {
 				testStores[j].add(alloc.randGen.Int63n(1<<20), 0)
 			}
-			if err := g.AddInfoProto(
+			if err := g.TestingAddInfoProtoAndWaitForAllCallbacks(
 				gossip.MakeStoreDescKey(roachpb.StoreID(j)),
 				&testStores[j].StoreDescriptor,
 				0,

pkg/testutils/gossiputil/BUILD.bazel

@@ -9,6 +9,5 @@ go_library(
     deps = [
         "//pkg/gossip",
         "//pkg/roachpb",
-        "//pkg/util/syncutil",
     ],
 )