feat: allow custom StreamSource heartbeat

Instead of using the default heartbeat for stream sources and mirror,
allow more customizability. The mechanism of the health check keeps
working in a linear fashion. Every time the heartbeat is not specified,
things work just like before.

Author: dodomorandi

server/stream.go

@@ -186,6 +186,7 @@ type StreamSource struct {
 	FilterSubject        string          `json:"filter_subject,omitempty"`
 	SubjectTransformDest string          `json:"subject_transform_dest,omitempty"`
 	External             *ExternalStream `json:"external,omitempty"`
+	Heartbeat            time.Duration   `json:"idle_heartbeat,omitempty"`
 
 	// Internal
 	iname string // For indexing when stream names are the same for multiple sources.
@@ -1898,7 +1899,7 @@ func (mset *stream) processMirrorMsgs(mirror *sourceInfo, ready *sync.WaitGroup)
 	// Signal the caller that we have captured the above fields.
 	ready.Done()
 
-	t := time.NewTicker(sourceHealthCheckInterval)
+	t := time.NewTicker(mset.getMirrorHeartbeat())
 	defer t.Stop()
 
 	for {
@@ -1920,7 +1921,7 @@ func (mset *stream) processMirrorMsgs(mirror *sourceInfo, ready *sync.WaitGroup)
 		case <-t.C:
 			mset.mu.RLock()
 			isLeader := mset.isLeader()
-			stalled := mset.mirror != nil && time.Since(mset.mirror.last) > 3*sourceHealthCheckInterval
+			stalled := mset.mirror != nil && time.Since(mset.mirror.last) > 3*mset.getMirrorHeartbeat()
 			mset.mu.RUnlock()
 			// No longer leader.
 			if !isLeader {
@@ -1937,6 +1938,14 @@ func (mset *stream) processMirrorMsgs(mirror *sourceInfo, ready *sync.WaitGroup)
 	}
 }
 
+func (mset *stream) getMirrorHeartbeat() time.Duration {
+	if mset.cfg.Mirror.Heartbeat == 0 {
+		return sourceHealthCheckInterval
+	} else {
+		return mset.cfg.Mirror.Heartbeat
+	}
+}
+
 // Checks that the message is from our current direct consumer. We can not depend on sub comparison
 // since cross account imports break.
 func (si *sourceInfo) isCurrentSub(reply string) bool {
@@ -2226,7 +2235,7 @@ func (mset *stream) setupMirrorConsumer() error {
 			AckPolicy:      AckNone,
 			AckWait:        22 * time.Hour,
 			MaxDeliver:     1,
-			Heartbeat:      sourceHealthCheckInterval,
+			Heartbeat:      mset.getMirrorHeartbeat(),
 			FlowControl:    true,
 			Direct:         true,
 		},
@@ -2537,7 +2546,7 @@ func (mset *stream) setSourceConsumer(iname string, seq uint64, startTime time.T
 			AckPolicy:      AckNone,
 			AckWait:        22 * time.Hour,
 			MaxDeliver:     1,
-			Heartbeat:      sourceHealthCheckInterval,
+			Heartbeat:      ssi.getStreamSourceHeartbeat(),
 			FlowControl:    true,
 			Direct:         true,
 		},
@@ -2712,12 +2721,13 @@ func (mset *stream) processSourceMsgs(si *sourceInfo, ready *sync.WaitGroup) {
 	msgs, qch, siqch, iname := si.msgs, mset.qch, si.qch, si.iname
 	// Set the last seen as now so that we don't fail at the first check.
 	si.last = time.Now()
+	heartbeat := mset.getSourceHeartbeat(iname)
 	mset.mu.Unlock()
 
 	// Signal the caller that we have captured the above fields.
 	ready.Done()
 
-	t := time.NewTicker(sourceHealthCheckInterval)
+	t := time.NewTicker(heartbeat)
 	defer t.Stop()
 
 	for {
@@ -2739,7 +2749,7 @@ func (mset *stream) processSourceMsgs(si *sourceInfo, ready *sync.WaitGroup) {
 		case <-t.C:
 			mset.mu.RLock()
 			isLeader := mset.isLeader()
-			stalled := time.Since(si.last) > 3*sourceHealthCheckInterval
+			stalled := time.Since(si.last) > 3*heartbeat
 			mset.mu.RUnlock()
 			// No longer leader.
 			if !isLeader {
@@ -5385,3 +5395,20 @@ func (mset *stream) clearMonitorRunning() {
 	defer mset.mu.Unlock()
 	mset.inMonitor = false
 }
+
+func (mset *stream) getSourceHeartbeat(iname string) time.Duration {
+	ssi := mset.streamSource(iname)
+	if ssi == nil {
+		return sourceHealthCheckInterval
+	}
+
+	return ssi.getStreamSourceHeartbeat()
+}
+
+func (ssi *StreamSource) getStreamSourceHeartbeat() time.Duration {
+	if ssi.Heartbeat == 0 {
+		return sourceHealthCheckInterval
+	} else {
+		return ssi.Heartbeat
+	}
+}