NRG: Periodic upper layer healthcheck

Signed-off-by: Neil Twigg <[email protected]>

Author: neilalexander

server/jetstream_cluster.go

@@ -955,6 +955,18 @@ func (js *jetStream) setupMetaGroup() error {
 		return err
 	}
 
+	// Quick healthcheck function to prove the metalayer *looks* functional.
+	n.SetUpperLayerHealthcheck(func() error {
+		// Check that we can acquire dios in a timely fashion.
+		select {
+		case <-dios:
+			dios <- struct{}{}
+		case <-time.After(3 * time.Second):
+			return fmt.Errorf("extreme dios contention")
+		}
+		return nil
+	})
+
 	// If we are bootstrapped with no state, start campaign early.
 	if bootstrap {
 		n.Campaign()
@@ -4354,11 +4366,15 @@ func (js *jetStream) processClusterUpdateStream(acc *Account, osa, sa *streamAss
 				mset.startClusterSubs()
 				mset.mu.Unlock()
 
-				js.createRaftGroup(acc.GetName(), rg, recovering, storage, pprofLabels{
+				// TODO(nat): Check error condition here, send SA result if err.
+				rn, _ := js.createRaftGroup(acc.GetName(), rg, recovering, storage, pprofLabels{
 					"type":    "stream",
 					"account": mset.accName(),
 					"stream":  mset.name(),
 				})
+				if rn != nil {
+					rn.SetUpperLayerHealthcheck(mset.upperLayerHealthcheck)
+				}
 			}
 			mset.monitorWg.Add(1)
 			// Start monitoring..
@@ -4545,11 +4561,14 @@ func (js *jetStream) processClusterCreateStream(acc *Account, sa *streamAssignme
 					s.Warnf("JetStream cluster error updating stream %q for account %q: %v", sa.Config.Name, acc.Name, err)
 					if osa != nil {
 						// Process the raft group and make sure it's running if needed.
-						js.createRaftGroup(acc.GetName(), osa.Group, osa.recovering, storage, pprofLabels{
+						rn, _ := js.createRaftGroup(acc.GetName(), osa.Group, osa.recovering, storage, pprofLabels{
 							"type":    "stream",
 							"account": mset.accName(),
 							"stream":  mset.name(),
 						})
+						if rn != nil {
+							rn.SetUpperLayerHealthcheck(mset.upperLayerHealthcheck)
+						}
 						mset.setStreamAssignment(osa)
 					}
 					if rg.node != nil {
@@ -5160,12 +5179,16 @@ func (js *jetStream) processClusterCreateConsumer(ca *consumerAssignment, state
 		storage = MemoryStorage
 	}
 	// No-op if R1.
-	js.createRaftGroup(accName, rg, recovering, storage, pprofLabels{
+	rn, _ := js.createRaftGroup(accName, rg, recovering, storage, pprofLabels{
 		"type":     "consumer",
 		"account":  mset.accName(),
 		"stream":   ca.Stream,
 		"consumer": ca.Name,
 	})
+	if rn != nil {
+		// For now consider a consumer to be healthy if the parent stream is.
+		rn.SetUpperLayerHealthcheck(mset.upperLayerHealthcheck)
+	}
 
 	// Check if we already have this consumer running.
 	var didCreate, isConfigUpdate, needsLocalResponse bool

server/raft.go

@@ -87,6 +87,7 @@ type RaftNode interface {
 	RecreateInternalSubs() error
 	IsSystemAccount() bool
 	GetTrafficAccountName() string
+	SetUpperLayerHealthcheck(f func() error)
 }
 
 type WAL interface {
@@ -233,6 +234,9 @@ type raft struct {
 	scaleUp      bool // The node is part of a scale up, puts us in observer mode until the log contains data.
 	membChanging bool // There is a membership change proposal in progress
 	deleted      bool // If the node was deleted.
+
+	hcf      func() error // Healthcheck function from the upper layer, optional
+	hcfailed bool         // Healthcheck function from the upper layer last reported failed
 }
 
 type proposedEntry struct {
@@ -625,6 +629,15 @@ func (n *raft) GetTrafficAccountName() string {
 	return n.acc.GetName()
 }
 
+// If provided, periodically check when leader if the upper layer is still responding.
+// If it does not return that it is happy then we will step down. The healthcheck
+// function can block, which is fine, because during that time runAsLeader can't send HBs.
+func (n *raft) SetUpperLayerHealthcheck(f func() error) {
+	n.Lock()
+	defer n.Unlock()
+	n.hcf = f
+}
+
 func (n *raft) RecreateInternalSubs() error {
 	n.Lock()
 	defer n.Unlock()
@@ -2219,6 +2232,15 @@ func (n *raft) processAppendEntries() {
 // runAsFollower is called by run and will block for as long as the node is
 // running in the follower state.
 func (n *raft) runAsFollower() {
+	// If we entered the follower state with a failed healthcheck then we
+	// will start a timer to keep checking it, as we might recover.
+	hct := make(<-chan time.Time)
+	if n.hcf != nil && n.hcfailed {
+		hcf := time.NewTicker(hbInterval * 30)
+		defer hcf.Stop()
+		hct = hcf.C
+	}
+
 	for n.State() == Follower {
 		elect := n.electTimer()
 
@@ -2254,6 +2276,12 @@ func (n *raft) runAsFollower() {
 				n.switchToCandidate()
 				return
 			}
+		case <-hct:
+			if err := n.hcf(); err == nil {
+				n.warn("Upper layer healthcheck failure has cleared")
+				n.setObserver(false, extUndetermined)
+				return
+			}
 		case <-n.votes.ch:
 			// We're receiving votes from the network, probably because we have only
 			// just stepped down and they were already in flight. Ignore them.
@@ -2793,6 +2821,9 @@ func (n *raft) runAsLeader() {
 	lq := time.NewTicker(lostQuorumCheck)
 	defer lq.Stop()
 
+	hcf := time.NewTicker(hbInterval * 30)
+	defer hcf.Stop()
+
 	for n.State() == Leader {
 		select {
 		case <-n.s.quitCh:
@@ -2850,6 +2881,18 @@ func (n *raft) runAsLeader() {
 				n.stepdown(noLeader)
 				return
 			}
+		case <-hcf.C:
+			if n.hcf == nil {
+				continue
+			}
+			if err := n.hcf(); err != nil {
+				n.hcfailed = true
+				n.warn("Upper layer healthcheck failed: %v", err)
+				n.warn("Stepping down from leader and switching to observer mode")
+				n.stepdown(noLeader)
+				n.setObserver(true, extUndetermined)
+				return
+			}
 		case <-n.votes.ch:
 			// Because of drain() it is possible that we get nil from popOne().
 			vresp, ok := n.votes.popOne()
@@ -3405,6 +3448,9 @@ func (n *raft) runAsCandidate() {
 	votes := map[string]struct{}{}
 	emptyVotes := map[string]struct{}{}
 
+	hcf := time.NewTicker(hbInterval * 30)
+	defer hcf.Stop()
+
 	for n.State() == Candidate {
 		elect := n.electTimer()
 		select {
@@ -3423,6 +3469,18 @@ func (n *raft) runAsCandidate() {
 		case <-elect.C:
 			n.switchToCandidate()
 			return
+		case <-hcf.C:
+			if n.hcf == nil {
+				continue
+			}
+			if err := n.hcf(); err != nil && !n.hcfailed {
+				n.hcfailed = true
+				n.warn("Upper layer healthcheck failed: %v", err)
+				n.warn("Stepping down from candidate and switching to observer mode")
+				n.stepdown(noLeader)
+				n.setObserver(true, extUndetermined)
+				return
+			}
 		case <-n.votes.ch:
 			// Because of drain() it is possible that we get nil from popOne().
 			vresp, ok := n.votes.popOne()
@@ -4759,6 +4817,17 @@ func (n *raft) switchToCandidate() {
 	n.Lock()
 	defer n.Unlock()
 
+	if n.hcf != nil {
+		if err := n.hcf(); err != nil {
+			n.debug("Not switching to candidate due to healthcheck failure:", err)
+			return
+		} else if n.hcfailed {
+			n.hcfailed = false
+			n.warn("Upper layer healthcheck failure has cleared")
+			n.setObserver(false, extUndetermined)
+		}
+	}
+
 	// If we are catching up or are in observer mode we can not switch.
 	// Avoid petitioning to become leader if we're behind on applies.
 	if n.observer || n.paused || n.processed < n.commit {

server/raft_test.go

@@ -4652,3 +4652,47 @@ func TestNRGMustNotResetVoteOnStepDownOrLeaderTransfer(t *testing.T) {
 	require_Equal(t, n.term, 1)
 	require_Equal(t, n.vote, nats0)
 }
+
+func TestNRGUpperLayerHealthcheck(t *testing.T) {
+	c := createJetStreamClusterExplicit(t, "R3S", 3)
+	defer c.shutdown()
+
+	rg := c.createMemRaftGroup("TEST", 3, newStateAdder)
+	leader := rg.waitOnLeader()
+	rn := leader.node()
+
+	start := time.Now()
+	rn.SetUpperLayerHealthcheck(func() error {
+		// Let's have normal operations for a bit before we signal
+		// an error.
+		if time.Since(start) > time.Second {
+			return fmt.Errorf("time has passed")
+		}
+		return nil
+	})
+
+	checkFor(t, 5*time.Second, 250*time.Millisecond, func() error {
+		rn := leader.node()
+		if rn.Leader() {
+			return fmt.Errorf("shouldn't still be leader")
+		}
+		if !rn.IsObserver() {
+			return fmt.Errorf("should be observer")
+		}
+		return nil
+	})
+
+	rn.SetUpperLayerHealthcheck(func() error {
+		// Now we'll set a healthcheck that always succeeds, so
+		// we should recover.
+		return nil
+	})
+
+	checkFor(t, 30*time.Second, 250*time.Millisecond, func() error {
+		rn := leader.node()
+		if rn.IsObserver() {
+			return fmt.Errorf("should not be observer")
+		}
+		return nil
+	})
+}

server/stream.go

@@ -1075,6 +1075,30 @@ func (mset *stream) setStreamAssignment(sa *streamAssignment) {
 	}
 }
 
+func (mset *stream) upperLayerHealthcheck() error {
+	// Check that nothing looks to be holding the lock for a long time.
+	// TODO(nat): we could check mset.mu etc here, but I'm not convinced we
+	// have eliminated all of the expensive or linear scans that may genuinely
+	// hold the lock and aren't worthy of a leadership change.
+	start := time.Now()
+	mset.cfgMu.RLock()
+	memory := mset.cfg.Storage == MemoryStorage
+	mset.cfgMu.RUnlock()
+	if time.Since(start) > 3*time.Second {
+		return fmt.Errorf("extreme stream lock contention")
+	}
+	if !memory {
+		// Check that we can acquire dios in a timely fashion.
+		select {
+		case <-dios:
+			dios <- struct{}{}
+		case <-time.After(3 * time.Second):
+			return fmt.Errorf("extreme dios contention")
+		}
+	}
+	return nil
+}
+
 func (mset *stream) monitorQuitC() <-chan struct{} {
 	if mset == nil {
 		return nil