Repository: nats-io/nats-server
Pull request #1965 (2021-03-05 13:49:42)

Push consumer idle heartbeats
Allow an option to push based consumers to have idle heartbeats delivered.
This allows an endpoint to know the consumer is still alive.
    
Signed-off-by: Derek Collison <derek@nats.io>

/cc @nats-io/core



Patch:
diff --git a/server/consumer.go b/server/consumer.go
index 1cd044a8..f54329e5 100644
--- a/server/consumer.go
+++ b/server/consumer.go
@@ -61,6 +61,7 @@ type ConsumerConfig struct {
 	SampleFrequency string        `json:"sample_freq,omitempty"`
 	MaxWaiting      int           `json:"max_waiting,omitempty"`
 	MaxAckPending   int           `json:"max_ack_pending,omitempty"`
+	Heartbeat       time.Duration `json:"idle_heartbeat,omitempty"`
 }
 
 type CreateConsumerRequest struct {
@@ -261,6 +262,9 @@ func (mset *stream) addConsumerWithAssignment(config *ConsumerConfig, oname stri
 		if config.MaxAckPending > 0 && config.AckPolicy == AckNone {
 			return nil, fmt.Errorf("consumer requires ack policy for max ack pending")
 		}
+		if config.Heartbeat > 0 && config.Heartbeat < 100*time.Millisecond {
+			return nil, fmt.Errorf("consumer idle heartbeat needs to be > 100ms")
+		}
 	} else {
 		// Pull mode / work queue mode require explicit ack.
 		if config.AckPolicy != AckExplicit {
@@ -281,6 +285,9 @@ func (mset *stream) addConsumerWithAssignment(config *ConsumerConfig, oname stri
 		if config.MaxWaiting == 0 {
 			config.MaxWaiting = JSWaitQueueDefaultMax
 		}
+		if config.Heartbeat > 0 {
+			return nil, fmt.Errorf("consumer idle heartbeat requires a push based consumer")
+		}
 	}
 
 	// Setup proper default for ack wait if we are in explicit ack mode.
@@ -1804,6 +1811,14 @@ func (o *consumer) checkWaitingForInterest() bool {
 	return o.waiting.len() > 0
 }
 
+// Lock should be held.
+func (o *consumer) hbTimer() (time.Duration, *time.Timer) {
+	if o.cfg.Heartbeat == 0 {
+		return 0, nil
+	}
+	return o.cfg.Heartbeat, time.NewTimer(o.cfg.Heartbeat)
+}
+
 func (o *consumer) loopAndGatherMsgs(qch chan struct{}) {
 	// On startup check to see if we are in a a reply situtation where replay policy is not instant.
 	var (
@@ -1820,6 +1835,11 @@ func (o *consumer) loopAndGatherMsgs(qch chan struct{}) {
 		}
 		lseq = o.mset.state().LastSeq
 	}
+	var hbc <-chan time.Time
+	hbd, hb := o.hbTimer()
+	if hb != nil {
+		hbc = hb.C
+	}
 	inch := o.inch
 	o.mu.Unlock()
 
@@ -1913,8 +1933,19 @@ func (o *consumer) loopAndGatherMsgs(qch chan struct{}) {
 			o.replay = false
 		}
 
+		// Reset our idle heartbeat timer if set.
+		if hb != nil {
+			if !hb.Stop() {
+				select {
+				case <-hbc:
+				default:
+				}
+			}
+			hb.Reset(hbd)
+		}
+
 		// We will wait here for new messages to arrive.
-		mch := o.mch
+		mch, outq, odsubj := o.mch, o.outq, o.cfg.DeliverSubject
 		o.mu.Unlock()
 
 		select {
@@ -1926,6 +1957,9 @@ func (o *consumer) loopAndGatherMsgs(qch chan struct{}) {
 			return
 		case <-mch:
 			// Messages are waiting.
+		case <-hbc:
+			hdr := []byte("NATS/1.0 100 Idle Heartbeat\r\n\r\n")
+			outq.send(&jsPubMsg{odsubj, odsubj, _EMPTY_, hdr, nil, nil, 0, nil})
 		}
 	}
 }
diff --git a/server/raft.go b/server/raft.go
index c0183854..c01a49f9 100644
--- a/server/raft.go
+++ b/server/raft.go
@@ -1254,8 +1254,11 @@ func (n *raft) resetElect(et time.Duration) {
 	if n.elect == nil {
 		n.elect = time.NewTimer(et)
 	} else {
-		if !n.elect.Stop() && len(n.elect.C) > 0 {
-			<-n.elect.C
+		if !n.elect.Stop() {
+			select {
+			case <-n.elect.C:
+			default:
+			}
 		}
 		n.elect.Reset(et)
 	}


Test patch:
diff --git a/server/jetstream_test.go b/server/jetstream_test.go
index ef2c3905..a017d443 100644
--- a/server/jetstream_test.go
+++ b/server/jetstream_test.go
@@ -7031,6 +7031,81 @@ func TestJetStreamSimpleFileRecovery(t *testing.T) {
 	}
 }
 
+func TestJetStreamPushConsumerIdleHeartbeats(t *testing.T) {
+	s := RunBasicJetStreamServer()
+	defer s.Shutdown()
+
+	// Forced cleanup of all persisted state.
+	if config := s.JetStreamConfig(); config != nil {
+		defer os.RemoveAll(config.StoreDir)
+	}
+
+	// Client for API requests.
+	nc, js := jsClientConnect(t, s)
+	defer nc.Close()
+
+	if _, err := js.AddStream(&nats.StreamConfig{Name: "TEST"}); err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+
+	sub, err := nc.SubscribeSync(nats.NewInbox())
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	defer sub.Unsubscribe()
+
+	// Test errors first
+	obsReq := CreateConsumerRequest{
+		Stream: "TEST",
+		Config: ConsumerConfig{
+			DeliverSubject: sub.Subject,
+			Heartbeat:      time.Millisecond,
+		},
+	}
+	req, err := json.Marshal(obsReq)
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	resp, err := nc.Request(fmt.Sprintf(JSApiConsumerCreateT, "TEST"), req, time.Second)
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	var ccResp JSApiConsumerCreateResponse
+	if err = json.Unmarshal(resp.Data, &ccResp); err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	if ccResp.Error == nil {
+		t.Fatalf("Expected an error, got none")
+	}
+	// Set acceptable heartbeat.
+	obsReq.Config.Heartbeat = 100 * time.Millisecond
+	req, err = json.Marshal(obsReq)
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	resp, err = nc.Request(fmt.Sprintf(JSApiConsumerCreateT, "TEST"), req, time.Second)
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	ccResp.Error = nil
+	if err = json.Unmarshal(resp.Data, &ccResp); err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	checkFor(t, time.Second, 20*time.Millisecond, func() error {
+		if nmsgs, _, err := sub.Pending(); err != nil || nmsgs < 9 {
+			return fmt.Errorf("Did not receive correct number of messages: %d vs %d", nmsgs, 9)
+		}
+		return nil
+	})
+	m, _ := sub.NextMsg(0)
+	if m.Header.Get("Status") != "100" {
+		t.Fatalf("Expected a 100 status code, got %q", m.Header.Get("Status"))
+	}
+	if m.Header.Get("Description") != "Idle Heartbeat" {
+		t.Fatalf("Wrong description , got %q", m.Header.Get("Description"))
+	}
+}
+
 func TestJetStreamInfoAPIWithHeaders(t *testing.T) {
 	s := RunBasicJetStreamServer()
 	defer s.Shutdown()
@@ -10484,6 +10559,9 @@ func TestJetStreamMirrorBasics(t *testing.T) {
 	// Now create our mirror stream.
 	createStreamOk(cfg)
 
+	// For now wait for the consumer state to register.
+	time.Sleep(250 * time.Millisecond)
+
 	// Send 100 messages.
 	for i := 0; i < 100; i++ {
 		if _, err := js.Publish("foo", []byte("OK")); err != nil {