smear heartbeats

Author: dodeca12

pkg/kv/kvserver/store_test.go

@@ -251,7 +251,7 @@ func createTestStoreWithoutStart(
 		supportGracePeriod := rpcContext.StoreLivenessWithdrawalGracePeriod()
 		options := storeliveness.NewOptions(livenessInterval, heartbeatInterval, supportGracePeriod)
 		transport, err := storeliveness.NewTransport(
-			cfg.AmbientCtx, stopper, cfg.Clock, cfg.NodeDialer, grpcServer, drpcServer, nil, /* knobs */
+			cfg.AmbientCtx, stopper, cfg.Clock, cfg.NodeDialer, grpcServer, drpcServer, cfg.Settings, nil, /* knobs */
 		)
 		require.NoError(t, err)
 		knobs := cfg.TestingKnobs.StoreLivenessKnobs

pkg/kv/kvserver/storeliveness/BUILD.bazel

@@ -34,6 +34,7 @@ go_library(
         "//pkg/util/protoutil",
         "//pkg/util/stop",
         "//pkg/util/syncutil",
+        "//pkg/util/taskpacer",
         "//pkg/util/timeutil",
         "@com_github_cockroachdb_errors//:errors",
         "@com_github_cockroachdb_redact//:redact",

pkg/kv/kvserver/storeliveness/support_manager.go

@@ -31,10 +31,38 @@ var Enabled = settings.RegisterBoolSetting(
 	true,
 )
 
+var UseHeartbeatCoordinator = settings.RegisterBoolSetting(
+	settings.SystemOnly,
+	"kv.store_liveness.use_heartbeat_coordinator",
+	"if enabled, heartbeat sends are the responsibility of the heartbeat "+
+		"coordinator which smears them over a certain duration, otherwise the "+
+		"heartbeat sends are sent immediately",
+	true,
+)
+
+var HeartbeatCoordinatorRefresh = settings.RegisterDurationSetting(
+	settings.SystemOnly,
+	"kv.store_liveness.heartbeat_coordinator.refresh",
+	"the refresh interval for the heartbeat coordinator pacer, which determines "+
+		"how frequently the coordinator checks for messages to send",
+	10*time.Millisecond,
+	settings.PositiveDuration,
+)
+
+var HeartbeatCoordinatorSmear = settings.RegisterDurationSetting(
+	settings.SystemOnly,
+	"kv.store_liveness.heartbeat_coordinator.smear",
+	"the smear duration for the heartbeat coordinator pacer, which spreads out "+
+		"send signals across multiple queues to avoid thundering herd",
+	1*time.Millisecond,
+	settings.PositiveDuration,
+)
+
 // MessageSender is the interface that defines how Store Liveness messages are
 // sent. Transport is the production implementation of MessageSender.
 type MessageSender interface {
 	EnqueueMessage(ctx context.Context, msg slpb.Message) (sent bool)
+	SendAllEnqueuedMessages(ctx context.Context)
 }
 
 // SupportManager orchestrates requesting and providing Store Liveness support.
@@ -168,6 +196,14 @@ func (sm *SupportManager) SupportFromEnabled(ctx context.Context) bool {
 	return Enabled.Get(&sm.settings.SV)
 }
 
+// SendViaHeartbeatCoordinator returns true if the cluster setting to send heartbeats
+// via the heartbeat coordinator (the goroutine in `transport.go`) is enabled.
+// The coordinator is responsible for smearing the heartbeats over
+// a certain duration, so expect heartbeats sends to be paced when enabled.
+func (sm *SupportManager) SendViaHeartbeatCoordinator(ctx context.Context) bool {
+	return UseHeartbeatCoordinator.Get(&sm.settings.SV)
+}
+
 // Start starts the main processing goroutine in startLoop as an async task.
 func (sm *SupportManager) Start(ctx context.Context) error {
 	// onRestart is called synchronously before the start of the main loop in
@@ -329,6 +365,9 @@ func (sm *SupportManager) sendHeartbeats(ctx context.Context) {
 			log.KvExec.Warningf(ctx, "failed to send heartbeat to store %+v", msg.To)
 		}
 	}
+	if sm.SendViaHeartbeatCoordinator(ctx) {
+		sm.sender.SendAllEnqueuedMessages(ctx)
+	}
 	sm.metrics.HeartbeatSuccesses.Inc(int64(successes))
 	sm.metrics.HeartbeatFailures.Inc(int64(len(heartbeats) - successes))
 	log.KvExec.VInfof(ctx, 2, "sent heartbeats to %d stores", successes)
@@ -426,6 +465,9 @@ func (sm *SupportManager) handleMessages(ctx context.Context, msgs []*slpb.Messa
 	for _, response := range responses {
 		_ = sm.sender.EnqueueMessage(ctx, response)
 	}
+	if sm.SendViaHeartbeatCoordinator(ctx) {
+		sm.sender.SendAllEnqueuedMessages(ctx)
+	}
 	log.KvExec.VInfof(ctx, 2, "sent %d heartbeat responses", len(responses))
 }
 

pkg/kv/kvserver/storeliveness/testutils.go

@@ -101,6 +101,10 @@ func (tms *testMessageSender) EnqueueMessage(_ context.Context, msg slpb.Message
 	return true
 }
 
+func (tms *testMessageSender) SendAllEnqueuedMessages(_ context.Context) {
+	// No-op for testing
+}
+
 func (tms *testMessageSender) drainSentMessages() []slpb.Message {
 	tms.mu.Lock()
 	defer tms.mu.Unlock()

pkg/kv/kvserver/storeliveness/transport.go

@@ -14,10 +14,12 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/rpc/nodedialer"
 	"github.com/cockroachdb/cockroach/pkg/rpc/rpcbase"
+	clustersettings "github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/stop"
 	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
+	"github.com/cockroachdb/cockroach/pkg/util/taskpacer"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/errors"
 	"google.golang.org/grpc"
@@ -54,7 +56,8 @@ type MessageHandler interface {
 
 // sendQueue is a queue of outgoing Messages.
 type sendQueue struct {
-	messages chan slpb.Message
+	sendMessages chan struct{}
+	messages     chan slpb.Message
 }
 
 // Transport handles the RPC messages for Store Liveness.
@@ -67,22 +70,44 @@ type sendQueue struct {
 // delivering them asynchronously.
 type Transport struct {
 	log.AmbientContext
-	stopper *stop.Stopper
-	clock   *hlc.Clock
-	dialer  *nodedialer.Dialer
-	metrics *TransportMetrics
+	stopper  *stop.Stopper
+	clock    *hlc.Clock
+	dialer   *nodedialer.Dialer
+	metrics  *TransportMetrics
+	settings *clustersettings.Settings
 
 	// queues stores outgoing message queues keyed by the destination node ID.
 	queues syncutil.Map[roachpb.NodeID, sendQueue]
 	// handlers stores the MessageHandler for each store on the node.
 	handlers syncutil.Map[roachpb.StoreID, MessageHandler]
 
+	// Once signaled, we will signal to all sendQueues to send all messages.
+	sendAllMessages chan struct{}
+
 	// TransportKnobs includes all knobs for testing.
 	knobs *TransportKnobs
 }
 
 var _ MessageSender = (*Transport)(nil)
 
+type pacerConfig struct {
+	settings *clustersettings.Settings
+}
+
+func (c pacerConfig) GetRefresh() time.Duration {
+	if c.settings != nil {
+		return HeartbeatCoordinatorRefresh.Get(&c.settings.SV)
+	}
+	return 10 * time.Millisecond
+}
+
+func (c pacerConfig) GetSmear() time.Duration {
+	if c.settings != nil {
+		return HeartbeatCoordinatorSmear.Get(&c.settings.SV)
+	}
+	return 1 * time.Millisecond
+}
+
 // NewTransport creates a new Store Liveness Transport.
 func NewTransport(
 	ambient log.AmbientContext,
@@ -91,18 +116,21 @@ func NewTransport(
 	dialer *nodedialer.Dialer,
 	grpcServer *grpc.Server,
 	drpcMux drpc.Mux,
+	settings *clustersettings.Settings,
 	knobs *TransportKnobs,
 ) (*Transport, error) {
 	if knobs == nil {
 		knobs = &TransportKnobs{}
 	}
 	t := &Transport{
-		AmbientContext: ambient,
-		stopper:        stopper,
-		clock:          clock,
-		dialer:         dialer,
-		metrics:        newTransportMetrics(),
-		knobs:          knobs,
+		AmbientContext:  ambient,
+		stopper:         stopper,
+		clock:           clock,
+		dialer:          dialer,
+		metrics:         newTransportMetrics(),
+		settings:        settings,
+		sendAllMessages: make(chan struct{}, 1),
+		knobs:           knobs,
 	}
 	if grpcServer != nil {
 		slpb.RegisterStoreLivenessServer(grpcServer, t)
@@ -112,6 +140,99 @@ func NewTransport(
 			return nil, err
 		}
 	}
+
+	// Start background goroutine to act as the transport sender coordinator. It
+	// is responsible for instructing the sendQueues to send their messages.
+	if err := stopper.RunAsyncTask(
+		context.Background(), "storeliveness transport send coordinator",
+		func(ctx context.Context) {
+			var batchTimer timeutil.Timer
+			defer batchTimer.Stop()
+
+			conf := pacerConfig{settings: settings}
+			pacer := taskpacer.New(conf)
+
+			// This will hold the channels we need to signal to send messages.
+			toSignal := make([]chan struct{}, 0)
+
+			for {
+				select {
+				case <-stopper.ShouldQuiesce():
+					return
+
+				case <-t.sendAllMessages:
+					// We received a signal to send all messages. Before we do that, let's
+					// wait for a short duration to give other stores a chance to
+					// enqueue messages which will increase batching opportunities.
+					batchTimer.Reset(batchDuration)
+					for done := false; !done; {
+						select {
+						case <-t.sendAllMessages:
+							// Consume any additional signals to send all messages.
+
+						case <-batchTimer.C:
+							// We have waited to batch messages
+							done = true
+						}
+					}
+
+					// At this point, we have waited for a short duration. We now need
+					// to signal all queues to send their messages.
+
+					// Get all the sendQueues that have messages to send. Note that the
+					// atomicity here is per sendQueue, and not across all sendQueues.
+					t.queues.Range(func(nodeID roachpb.NodeID, q *sendQueue) bool {
+						if len(q.messages) == 0 {
+							// Nothing to send.
+							return true
+						}
+
+						toSignal = append(toSignal, q.sendMessages)
+						return true
+					})
+
+					// There is a benign race condition here, and it happens in two cases:
+					// 1. If after we inserted the toSignal channels, a new message is
+					// enqueued to a new queue that we haven't added. In this case, the
+					// t.sendAllMessages should be set, and we will pick it up in the next
+					// iteration of the for loop.
+					// 2. If after we inserted the toSignal channels, a new message is
+					// added to a queue that we have already added. In this case, in the
+					// next iteration t.sendAllMessages might be valid, but the queues
+					// could be empty. This is not a problem because we won't wake up
+					// any sendQueue goroutine unnecessarily.
+
+					// Pace the signaling of the channels.
+					pacer.StartTask(timeutil.Now())
+					workLeft := len(toSignal)
+					for workLeft > 0 {
+						todo, by := pacer.Pace(timeutil.Now(), workLeft)
+
+						// Pop todo items off the toSignal slice and signal them.
+						for i := 0; i < todo && workLeft > 0; i++ {
+							ch := toSignal[len(toSignal)-1]
+							toSignal = toSignal[:len(toSignal)-1]
+							select {
+							case ch <- struct{}{}:
+							default:
+							}
+							workLeft--
+						}
+
+						if workLeft > 0 && timeutil.Now().Before(by) {
+							time.Sleep(by.Sub(timeutil.Now()))
+						}
+					}
+
+					// Clear toSignal for next iteration
+					toSignal = toSignal[:0]
+				}
+			}
+		},
+	); err != nil {
+		return nil, err
+	}
+
 	return t, nil
 }
 
@@ -265,12 +386,23 @@ func (t *Transport) EnqueueMessage(ctx context.Context, msg slpb.Message) (enque
 	}
 }
 
+// SendAllEnqueuedMessages signals all queues to send all their messages.
+func (t *Transport) SendAllEnqueuedMessages(ctx context.Context) {
+	select {
+	case t.sendAllMessages <- struct{}{}:
+	default:
+	}
+}
+
 // getQueue returns the queue for the specified node ID and a boolean
 // indicating whether the queue already exists (true) or was created (false).
 func (t *Transport) getQueue(nodeID roachpb.NodeID) (*sendQueue, bool) {
 	queue, ok := t.queues.Load(nodeID)
 	if !ok {
-		q := sendQueue{messages: make(chan slpb.Message, sendBufferSize)}
+		q := sendQueue{
+			sendMessages: make(chan struct{}, 1),
+			messages:     make(chan slpb.Message, sendBufferSize),
+		}
 		queue, ok = t.queues.LoadOrStore(nodeID, &q)
 	}
 	return queue, ok
@@ -365,6 +497,53 @@ func (t *Transport) processQueue(
 	var batchTimer timeutil.Timer
 	defer batchTimer.Stop()
 	batch := &slpb.MessageBatch{}
+
+	drainQueue := func() {
+		for {
+			select {
+			case msg := <-q.messages:
+				batch.Messages = append(batch.Messages, msg)
+				t.metrics.SendQueueSize.Dec(1)
+				t.metrics.SendQueueBytes.Dec(int64(msg.Size()))
+			default:
+				return
+			}
+		}
+	}
+
+	sendBatch := func() error {
+		if len(batch.Messages) == 0 {
+			return nil
+		}
+		batch.Now = t.clock.NowAsClockTimestamp()
+		if err := stream.Send(batch); err != nil {
+			t.metrics.MessagesSendDropped.Inc(int64(len(batch.Messages)))
+			return err
+		}
+		t.metrics.BatchesSent.Inc(1)
+		t.metrics.MessagesSent.Inc(int64(len(batch.Messages)))
+		// Reuse the Messages slice, but zero out the contents to avoid delaying GC.
+		for i := range batch.Messages {
+			batch.Messages[i] = slpb.Message{}
+		}
+		batch.Messages = batch.Messages[:0]
+		batch.Now = hlc.ClockTimestamp{}
+		return nil
+	}
+
+	// Check if we should use the heartbeat coordinator
+	var useHeartbeatCoordinator bool
+	if t.settings != nil {
+		useHeartbeatCoordinator = UseHeartbeatCoordinator.Get(&t.settings.SV)
+	}
+
+	// If coordinator is enabled, we don't want to process messages directly.
+	// We'll set the messages channel to nil so the select never triggers on it.
+	var directMessages <-chan slpb.Message
+	if !useHeartbeatCoordinator {
+		directMessages = q.messages
+	}
+
 	for {
 		idleTimer.Reset(getIdleTimeout())
 		select {
@@ -375,7 +554,13 @@ func (t *Transport) processQueue(
 			t.metrics.SendQueueIdle.Inc(1)
 			return nil
 
-		case msg := <-q.messages:
+		case <-q.sendMessages:
+			drainQueue()
+			if err = sendBatch(); err != nil {
+				return err
+			}
+
+		case msg := <-directMessages:
 			batch.Messages = append(batch.Messages, msg)
 			t.metrics.SendQueueSize.Dec(1)
 			t.metrics.SendQueueBytes.Dec(int64(msg.Size()))