Repository: nats-io/nats.go
Pull request #1402 (2023-09-15 11:01:50)

[ADDED] Setting heartbeats in Fetch and FetchBatch in legacy API


Patch:
diff --git a/js.go b/js.go
index a8deb98..ae19641 100644
--- a/js.go
+++ b/js.go
@@ -1176,10 +1176,11 @@ type SequencePair struct {
 
 // nextRequest is for getting next messages for pull based consumers from JetStream.
 type nextRequest struct {
-	Expires  time.Duration `json:"expires,omitempty"`
-	Batch    int           `json:"batch,omitempty"`
-	NoWait   bool          `json:"no_wait,omitempty"`
-	MaxBytes int           `json:"max_bytes,omitempty"`
+	Expires   time.Duration `json:"expires,omitempty"`
+	Batch     int           `json:"batch,omitempty"`
+	NoWait    bool          `json:"no_wait,omitempty"`
+	MaxBytes  int           `json:"max_bytes,omitempty"`
+	Heartbeat time.Duration `json:"idle_heartbeat,omitempty"`
 }
 
 // jsSub includes JetStream subscription info.
@@ -2469,6 +2470,7 @@ func EnableFlowControl() SubOpt {
 }
 
 // IdleHeartbeat enables push based consumers to have idle heartbeats delivered.
+// For pull consumers, idle heartbeat has to be set on each [Fetch] call.
 func IdleHeartbeat(duration time.Duration) SubOpt {
 	return subOptFn(func(opts *subOpts) error {
 		opts.cfg.Heartbeat = duration
@@ -2588,6 +2590,7 @@ type pullOpts struct {
 	maxBytes int
 	ttl      time.Duration
 	ctx      context.Context
+	hb       time.Duration
 }
 
 // PullOpt are the options that can be passed when pulling a batch of messages.
@@ -2603,6 +2606,16 @@ func PullMaxWaiting(n int) SubOpt {
 	})
 }
 
+type PullHeartbeat time.Duration
+
+func (h PullHeartbeat) configurePull(opts *pullOpts) error {
+	if h <= 0 {
+		return fmt.Errorf("%w: idle heartbeat has to be greater than 0", ErrInvalidArg)
+	}
+	opts.hb = time.Duration(h)
+	return nil
+}
+
 // PullMaxBytes defines the max bytes allowed for a fetch request.
 type PullMaxBytes int
 
@@ -2646,6 +2659,11 @@ func checkMsg(msg *Msg, checkSts, isNoWait bool) (usrMsg bool, err error) {
 	if !checkSts {
 		return
 	}
+
+	// if it's a heartbeat message, report as not user msg
+	if isHb, _ := isJSControlMessage(msg); isHb {
+		return
+	}
 	switch val {
 	case noResponders:
 		err = ErrNoResponders
@@ -2722,6 +2740,9 @@ func (sub *Subscription) Fetch(batch int, opts ...PullOpt) ([]*Msg, error) {
 		ttl = js.opts.wait
 	}
 	sub.mu.Unlock()
+	if o.hb != 0 && 2*o.hb >= ttl {
+		return nil, fmt.Errorf("%w: idle heartbeat value too large", ErrInvalidArg)
+	}
 
 	// Use the given context or setup a default one for the span
 	// of the pull batch request.
@@ -2732,7 +2753,6 @@ func (sub *Subscription) Fetch(batch int, opts ...PullOpt) ([]*Msg, error) {
 	)
 	if ctx == nil {
 		ctx, cancel = context.WithTimeout(context.Background(), ttl)
-		defer cancel()
 	} else if _, hasDeadline := ctx.Deadline(); !hasDeadline {
 		// Prevent from passing the background context which will just block
 		// and cannot be canceled either.
@@ -2743,8 +2763,10 @@ func (sub *Subscription) Fetch(batch int, opts ...PullOpt) ([]*Msg, error) {
 		// If the context did not have a deadline, then create a new child context
 		// that will use the default timeout from the JS context.
 		ctx, cancel = context.WithTimeout(ctx, ttl)
-		defer cancel()
+	} else {
+		ctx, cancel = context.WithCancel(ctx)
 	}
+	defer cancel()
 
 	// Check if context not done already before making the request.
 	select {
@@ -2783,6 +2805,8 @@ func (sub *Subscription) Fetch(batch int, opts ...PullOpt) ([]*Msg, error) {
 			msgs = append(msgs, msg)
 		}
 	}
+	var hbTimer *time.Timer
+	var hbErr error
 	if err == nil && len(msgs) < batch {
 		// For batch real size of 1, it does not make sense to set no_wait in
 		// the request.
@@ -2813,8 +2837,26 @@ func (sub *Subscription) Fetch(batch int, opts ...PullOpt) ([]*Msg, error) {
 			nr.Expires = expires
 			nr.NoWait = noWait
 			nr.MaxBytes = o.maxBytes
+			if 2*o.hb < expires {
+				nr.Heartbeat = o.hb
+			} else {
+				nr.Heartbeat = 0
+			}
 			req, _ := json.Marshal(nr)
-			return nc.PublishRequest(nms, rply, req)
+			if err := nc.PublishRequest(nms, rply, req); err != nil {
+				return err
+			}
+			if o.hb > 0 {
+				if hbTimer == nil {
+					hbTimer = time.AfterFunc(2*o.hb, func() {
+						hbErr = ErrNoHeartbeat
+						cancel()
+					})
+				} else {
+					hbTimer.Reset(2 * o.hb)
+				}
+			}
+			return nil
 		}
 
 		err = sendReq()
@@ -2822,6 +2864,9 @@ func (sub *Subscription) Fetch(batch int, opts ...PullOpt) ([]*Msg, error) {
 			// Ask for next message and wait if there are no messages
 			msg, err = sub.nextMsgWithContext(ctx, true, true)
 			if err == nil {
+				if hbTimer != nil {
+					hbTimer.Reset(2 * o.hb)
+				}
 				var usrMsg bool
 
 				usrMsg, err = checkMsg(msg, true, noWait)
@@ -2840,9 +2885,15 @@ func (sub *Subscription) Fetch(batch int, opts ...PullOpt) ([]*Msg, error) {
 				}
 			}
 		}
+		if hbTimer != nil {
+			hbTimer.Stop()
+		}
 	}
 	// If there is at least a message added to msgs, then need to return OK and no error
 	if err != nil && len(msgs) == 0 {
+		if hbErr != nil {
+			return nil, hbErr
+		}
 		return nil, o.checkCtxErr(err)
 	}
 	return msgs, nil
@@ -2950,6 +3001,9 @@ func (sub *Subscription) FetchBatch(batch int, opts ...PullOpt) (MessageBatch, e
 		ttl = js.opts.wait
 	}
 	sub.mu.Unlock()
+	if o.hb != 0 && 2*o.hb >= ttl {
+		return nil, fmt.Errorf("%w: idle heartbeat value too large", ErrInvalidArg)
+	}
 
 	// Use the given context or setup a default one for the span
 	// of the pull batch request.
@@ -2970,10 +3024,12 @@ func (sub *Subscription) FetchBatch(batch int, opts ...PullOpt) (MessageBatch, e
 		// If the context did not have a deadline, then create a new child context
 		// that will use the default timeout from the JS context.
 		ctx, cancel = context.WithTimeout(ctx, ttl)
+	} else {
+		ctx, cancel = context.WithCancel(ctx)
 	}
 	defer func() {
 		// only cancel the context here if we are sure the fetching goroutine has not been started yet
-		if cancel != nil && cancelContext {
+		if cancelContext {
 			cancel()
 		}
 	}()
@@ -3031,9 +3087,10 @@ func (sub *Subscription) FetchBatch(batch int, opts ...PullOpt) (MessageBatch, e
 
 	requestBatch := batch - len(result.msgs)
 	req := nextRequest{
-		Expires:  expires,
-		Batch:    requestBatch,
-		MaxBytes: o.maxBytes,
+		Expires:   expires,
+		Batch:     requestBatch,
+		MaxBytes:  o.maxBytes,
+		Heartbeat: o.hb,
 	}
 	reqJSON, err := json.Marshal(req)
 	if err != nil {
@@ -3051,11 +3108,17 @@ func (sub *Subscription) FetchBatch(batch int, opts ...PullOpt) (MessageBatch, e
 		result.err = err
 		return result, nil
 	}
+	var hbTimer *time.Timer
+	var hbErr error
+	if o.hb > 0 {
+		hbTimer = time.AfterFunc(2*o.hb, func() {
+			hbErr = ErrNoHeartbeat
+			cancel()
+		})
+	}
 	cancelContext = false
 	go func() {
-		if cancel != nil {
-			defer cancel()
-		}
+		defer cancel()
 		var requestMsgs int
 		for requestMsgs < requestBatch {
 			// Ask for next message and wait if there are no messages
@@ -3063,6 +3126,9 @@ func (sub *Subscription) FetchBatch(batch int, opts ...PullOpt) (MessageBatch, e
 			if err != nil {
 				break
 			}
+			if hbTimer != nil {
+				hbTimer.Reset(2 * o.hb)
+			}
 			var usrMsg bool
 
 			usrMsg, err = checkMsg(msg, true, false)
@@ -3082,7 +3148,11 @@ func (sub *Subscription) FetchBatch(batch int, opts ...PullOpt) (MessageBatch, e
 			}
 		}
 		if err != nil {
-			result.err = o.checkCtxErr(err)
+			if hbErr != nil {
+				result.err = hbErr
+			} else {
+				result.err = o.checkCtxErr(err)
+			}
 		}
 		close(result.msgs)
 		result.done <- struct{}{}
diff --git a/jserrors.go b/jserrors.go
index d7959ca..cb3845c 100644
--- a/jserrors.go
+++ b/jserrors.go
@@ -104,6 +104,9 @@ var (
 	// ErrConsumerLeadershipChanged is returned when pending requests are no longer valid after leadership has changed
 	ErrConsumerLeadershipChanged JetStreamError = &jsError{message: "Leadership Changed"}
 
+	// ErrNoHeartbeat is returned when no heartbeat is received from server when sending requests with pull consumer.
+	ErrNoHeartbeat JetStreamError = &jsError{message: "no heartbeat received"}
+
 	// DEPRECATED: ErrInvalidDurableName is no longer returned and will be removed in future releases.
 	// Use ErrInvalidConsumerName instead.
 	ErrInvalidDurableName = errors.New("nats: invalid durable name")
diff --git a/nats.go b/nats.go
index 36fb9df..821e823 100644
--- a/nats.go
+++ b/nats.go
@@ -3168,8 +3168,10 @@ func (nc *Conn) processMsg(data []byte) {
 		}
 	}
 
-	// Skip processing if this is a control message.
-	if !ctrlMsg {
+	// Skip processing if this is a control message and
+	// if not a pull consumer heartbeat. For pull consumers,
+	// heartbeats have to be handled on per request basis.
+	if !ctrlMsg || (jsi != nil && jsi.pull) {
 		var chanSubCheckFC bool
 		// Subscription internal stats (applicable only for non ChanSubscription's)
 		if sub.typ != ChanSubscription {


Test patch:
diff --git a/test/js_test.go b/test/js_test.go
index d9ad8ca..ab599c4 100644
--- a/test/js_test.go
+++ b/test/js_test.go
@@ -1124,6 +1124,176 @@ func TestJetStreamSubscribe_SkipConsumerLookup(t *testing.T) {
 	})
 }
 
+func TestPullSubscribeFetchWithHeartbeat(t *testing.T) {
+	s := RunBasicJetStreamServer()
+	defer shutdownJSServerAndRemoveStorage(t, s)
+
+	nc, js := jsClient(t, s)
+	defer nc.Close()
+
+	_, err := js.AddStream(&nats.StreamConfig{
+		Name:     "TEST",
+		Subjects: []string{"foo"},
+	})
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+
+	sub, err := js.PullSubscribe("foo", "")
+	if err != nil {
+		t.Fatalf("Unexpected error: %s", err)
+	}
+	for i := 0; i < 5; i++ {
+		if _, err := js.Publish("foo", []byte("msg")); err != nil {
+			t.Fatalf("Unexpected error: %s", err)
+		}
+	}
+
+	// fetch 5 messages, should finish immediately
+	msgs, err := sub.Fetch(5, nats.PullHeartbeat(100*time.Millisecond))
+	if err != nil {
+		t.Fatalf("Unexpected error: %s", err)
+	}
+	if len(msgs) != 5 {
+		t.Fatalf("Expected %d messages; got: %d", 5, len(msgs))
+	}
+	now := time.Now()
+	// no messages available, should time out normally
+	_, err = sub.Fetch(5, nats.PullHeartbeat(50*time.Millisecond), nats.MaxWait(300*time.Millisecond))
+	elapsed := time.Since(now)
+	if elapsed < 300*time.Millisecond {
+		t.Fatalf("Expected timeout after 300ms; got: %v", elapsed)
+	}
+	if !errors.Is(err, nats.ErrTimeout) {
+		t.Fatalf("Expected timeout error; got: %v", err)
+	}
+
+	// delete consumer to verify heartbeats are not sent anymore
+	info, err := sub.ConsumerInfo()
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	if err := js.DeleteConsumer("TEST", info.Name); err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	_, err = sub.Fetch(5, nats.PullHeartbeat(100*time.Millisecond), nats.MaxWait(1*time.Second))
+	if !errors.Is(err, nats.ErrNoHeartbeat) {
+		t.Fatalf("Expected no heartbeat error; got: %v", err)
+	}
+
+	// heartbeat value too large
+	_, err = sub.Fetch(5, nats.PullHeartbeat(200*time.Millisecond), nats.MaxWait(300*time.Millisecond))
+	if !errors.Is(err, nats.ErrInvalidArg) {
+		t.Fatalf("Expected no heartbeat error; got: %v", err)
+	}
+
+	// heartbeat value invalid
+	_, err = sub.Fetch(5, nats.PullHeartbeat(-1))
+	if !errors.Is(err, nats.ErrInvalidArg) {
+		t.Fatalf("Expected no heartbeat error; got: %v", err)
+	}
+}
+
+func TestPullSubscribeFetchBatchWithHeartbeat(t *testing.T) {
+	s := RunBasicJetStreamServer()
+	defer shutdownJSServerAndRemoveStorage(t, s)
+
+	nc, js := jsClient(t, s)
+	defer nc.Close()
+
+	_, err := js.AddStream(&nats.StreamConfig{
+		Name:     "TEST",
+		Subjects: []string{"foo"},
+	})
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+
+	sub, err := js.PullSubscribe("foo", "")
+	if err != nil {
+		t.Fatalf("Unexpected error: %s", err)
+	}
+	for i := 0; i < 5; i++ {
+		if _, err := js.Publish("foo", []byte("msg")); err != nil {
+			t.Fatalf("Unexpected error: %s", err)
+		}
+	}
+
+	// fetch 5 messages, should finish immediately
+	msgs, err := sub.FetchBatch(5, nats.PullHeartbeat(100*time.Millisecond))
+	if err != nil {
+		t.Fatalf("Unexpected error: %s", err)
+	}
+	var i int
+	for msg := range msgs.Messages() {
+		i++
+		msg.Ack()
+	}
+	if i != 5 {
+		t.Fatalf("Expected %d messages; got: %d", 5, i)
+	}
+	if msgs.Error() != nil {
+		t.Fatalf("Unexpected error: %s", msgs.Error())
+	}
+	now := time.Now()
+	// no messages available, should time out normally
+	msgs, err = sub.FetchBatch(5, nats.PullHeartbeat(50*time.Millisecond), nats.MaxWait(300*time.Millisecond))
+	if err != nil {
+		t.Fatalf("Unexpected error: %s", err)
+	}
+	i = 0
+	for msg := range msgs.Messages() {
+		i++
+		msg.Ack()
+	}
+	elapsed := time.Since(now)
+	if i != 0 {
+		t.Fatalf("Expected %d messages; got: %d", 0, i)
+	}
+	if msgs.Error() != nil {
+		t.Fatalf("Unexpected error: %s", msgs.Error())
+	}
+	if elapsed < 290*time.Millisecond {
+		t.Fatalf("Expected timeout after 300ms; got: %v", elapsed)
+	}
+
+	// delete consumer to verify heartbeats are not sent anymore
+	info, err := sub.ConsumerInfo()
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	if err := js.DeleteConsumer("TEST", info.Name); err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	now = time.Now()
+	msgs, err = sub.FetchBatch(5, nats.PullHeartbeat(100*time.Millisecond), nats.MaxWait(1*time.Second))
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	for msg := range msgs.Messages() {
+		msg.Ack()
+	}
+	elapsed = time.Since(now)
+	if elapsed < 200*time.Millisecond || elapsed > 300*time.Millisecond {
+		t.Fatalf("Expected timeout after 200ms and before 300ms; got: %v", elapsed)
+	}
+	if !errors.Is(msgs.Error(), nats.ErrNoHeartbeat) {
+		t.Fatalf("Expected no heartbeat error; got: %v", err)
+	}
+
+	// heartbeat value too large
+	_, err = sub.Fetch(5, nats.PullHeartbeat(200*time.Millisecond), nats.MaxWait(300*time.Millisecond))
+	if !errors.Is(err, nats.ErrInvalidArg) {
+		t.Fatalf("Expected no heartbeat error; got: %v", err)
+	}
+
+	// heartbeat value invalid
+	_, err = sub.Fetch(5, nats.PullHeartbeat(-1))
+	if !errors.Is(err, nats.ErrInvalidArg) {
+		t.Fatalf("Expected no heartbeat error; got: %v", err)
+	}
+}
+
 func TestPullSubscribeFetchBatch(t *testing.T) {
 	s := RunBasicJetStreamServer()
 	defer shutdownJSServerAndRemoveStorage(t, s)