Merge pull request #152467 from iskettaneh/backport25.2-144557

release-25.2: gossip: parallelize callback handling

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/gossip_test.go

@@ -1140,3 +1140,70 @@ func TestGossipBatching(t *testing.T) {
 	require.LessOrEqual(t, serverMessagesSentCount, upperBoundMessages)
 	require.LessOrEqual(t, clientMessagesSentCount, upperBoundMessages)
 }
+
+// TestCallbacksPendingMetricGoesToZeroOnStop verifies that the CallbacksPending
+// metric is correctly decremented when a callback is unregistered with pending work
+// or when the stopper is stopped.
+func TestCallbacksPendingMetricGoesToZeroOnStop(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+
+	testCases := []struct {
+		name    string
+		cleanup func(g *Gossip, unregister func(), stopper *stop.Stopper, ctx context.Context)
+	}{
+		{
+			name: "unregister callback",
+			cleanup: func(g *Gossip, unregister func(), stopper *stop.Stopper, ctx context.Context) {
+				unregister()
+			},
+		},
+		{
+			name: "stopper shutdown",
+			cleanup: func(g *Gossip, unregister func(), stopper *stop.Stopper, ctx context.Context) {
+				stopper.Stop(ctx)
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			ctx := context.Background()
+			stopper := stop.NewStopper()
+			defer stopper.Stop(ctx)
+			g := NewTest(1, stopper, metric.NewRegistry())
+
+			unregister := g.RegisterCallback("test.*", func(key string, val roachpb.Value) {
+				// Do nothing.
+			})
+
+			// Add 100 infos to the gossip that will be processed by the callback.
+			for i := 0; i < 100; i++ {
+				slice := []byte("b1")
+				require.NoError(t, g.AddInfo(fmt.Sprintf("test.key%d", i), slice, time.Hour))
+			}
+
+			// Execute the cleanup action (either unregister or stopper.Stop)
+			// We do this in a goroutine to help cause interesting potential race conditions.
+			go func() {
+				tc.cleanup(g, unregister, stopper, ctx)
+			}()
+
+			// Add another 100 infos to the gossip that will be processed by the callback.
+			// We do this in a goroutine to help cause interesting potential race conditions.
+			go func() {
+				for i := 0; i < 100; i++ {
+					slice := []byte("b2")
+					require.NoError(t, g.AddInfo(fmt.Sprintf("test.key%d", i), slice, time.Hour))
+				}
+			}()
+
+			// Wait for the pending callbacks metric to go to 0.
+			testutils.SucceedsSoon(t, func() error {
+				if g.mu.is.metrics.CallbacksPending.Value() != 0 {
+					return fmt.Errorf("CallbacksPending should be 0, got %d", g.mu.is.metrics.CallbacksPending.Value())
+				}
+				return nil
+			})
+		})
+	}
+}