Check for shutdown in heartbeat fast-path

raft.go

@@ -230,7 +230,7 @@ func NewRaft(conf *Config, fsm FSM, logs LogStore, stable StableStore, snaps Sna
 	// Setup a heartbeat fast-path to avoid head-of-line
 	// blocking where possible. It MUST be safe for this
 	// to be called concurrently with a blocking RPC.
-	trans.SetHeartbeatHandler(r.processRPC)
+	trans.SetHeartbeatHandler(r.processHeartbeat)
 
 	// Start the background work
 	r.goFunc(r.run)
@@ -1178,6 +1178,26 @@ func (r *Raft) processRPC(rpc RPC) {
 	}
 }
 
+// processHeartbeat is a special handler used just for heartbeat requests
+// so that they can be fast-pathed if a transport supports it
+func (r *Raft) processHeartbeat(rpc RPC) {
+	// Check if we are shutdown, just ignore the RPC
+	select {
+	case <-r.shutdownCh:
+		return
+	default:
+	}
+
+	// Ensure we are only handling a heartbeat
+	switch cmd := rpc.Command.(type) {
+	case *AppendEntriesRequest:
+		r.appendEntries(rpc, cmd)
+	default:
+		r.logger.Printf("[ERR] raft: Expected heartbeat, got command: %#v", rpc.Command)
+		rpc.Respond(nil, fmt.Errorf("unexpected command"))
+	}
+}
+
 // appendEntries is invoked when we get an append entries RPC call
 func (r *Raft) appendEntries(rpc RPC, a *AppendEntriesRequest) {
 	defer metrics.MeasureSince([]string{"raft", "rpc", "appendEntries"}, time.Now())