Request idle heartbeats for ordered consumers

nats-jetstream/src/nats/jetstream/consumer/ordered.py

@@ -27,6 +27,11 @@
 
 logger = logging.getLogger(__name__)
 
+# How long a pull request rests on the server before expiring, and how often the
+# server reports it is still alive while it waits.
+_DEFAULT_MAX_WAIT = 30.0
+_DEFAULT_HEARTBEAT = 15.0
+
 
 @dataclass
 class _Cursor:
@@ -457,9 +462,14 @@ async def _create_inner_stream(self) -> PullMessageStream:
         """Create a new inner PullMessageStream from the current consumer."""
         if self._consumer._current_consumer is None:
             raise OrderedConsumerClosedError("No active consumer")
+        max_wait = self._max_wait if self._max_wait is not None else _DEFAULT_MAX_WAIT
         stream = await self._consumer._current_consumer.messages(
-            heartbeat=self._heartbeat,
-            max_wait=self._max_wait if self._max_wait is not None else 30.0,
+            # An ordered consumer resets when its inner stream ends, and without
+            # idle heartbeats nothing ends it: a consumer that stops answering
+            # leaves the iterator waiting forever. Always ask for heartbeats,
+            # keeping them under the request expiry the server enforces.
+            heartbeat=self._heartbeat if self._heartbeat is not None else min(_DEFAULT_HEARTBEAT, max_wait / 2),
+            max_wait=max_wait,
             max_messages=self._max_messages,
             max_bytes=self._max_bytes,
         )

nats-jetstream/src/nats/jetstream/consumer/pull.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+import contextlib
 import json
 import logging
 import time
@@ -31,6 +32,10 @@
 
 logger = logging.getLogger(__name__)
 
+# Consecutive missed heartbeat windows before a stream is considered dead. One
+# miss can be a lost batch; a second means nothing is answering.
+_MAX_MISSED_HEARTBEATS = 2
+
 
 def _message_size(msg: ClientMessage) -> int:
     """Calculate the size of a message per ADR-37.
@@ -233,6 +238,7 @@ class PullMessageStream(MessageStream):
     _heartbeat_deadline: float | None
     _heartbeat_paused: bool
     _heartbeat_remaining: float | None
+    _missed_heartbeats: int
 
     def __init__(
         self,
@@ -275,6 +281,7 @@ def __init__(
         self._heartbeat_task: asyncio.Task | None = None
         self._started = False
         self._heartbeat_deadline = time.time() + (heartbeat * 2) if heartbeat is not None else None
+        self._missed_heartbeats = 0
 
         # Register disconnect/reconnect callbacks for heartbeat timer (ADR-37)
         if heartbeat is not None:
@@ -326,6 +333,7 @@ async def __anext__(self) -> Message:
             # Reset heartbeat timer on any message (ADR-37)
             if self._heartbeat is not None:
                 self._heartbeat_deadline = time.time() + (self._heartbeat * 2)
+                self._missed_heartbeats = 0
 
             # Handle status messages
             if raw_msg.status is not None:
@@ -455,12 +463,29 @@ async def _heartbeat_monitor(self):
 
                 # Check if heartbeat timeout has been reached (2x idle_heartbeat)
                 if self._heartbeat_deadline is not None and time.time() > self._heartbeat_deadline:
+                    self._missed_heartbeats += 1
                     logger.warning(
-                        "Heartbeat timeout: no message received within %.2fs (2x idle_heartbeat of %.2fs)",
+                        "Heartbeat timeout %d: no message received within %.2fs (2x idle_heartbeat of %.2fs)",
+                        self._missed_heartbeats,
                         self._heartbeat * 2,
                         self._heartbeat,
                     )
-                    # Reset pending counts and request more messages (non-terminal)
+
+                    if self._missed_heartbeats >= _MAX_MISSED_HEARTBEATS:
+                        # The consumer is unreachable -- it may have been deleted
+                        # or lost with the server it lived on. End the stream so
+                        # the caller sees it; an ordered consumer takes this as
+                        # its cue to reset and recreate. Unsubscribing wakes a
+                        # pending __anext__, which finishes cleanup; this task
+                        # must not run _cleanup itself because that would cancel
+                        # and await the task it is running on.
+                        self._terminated = True
+                        with contextlib.suppress(Exception):
+                            await self._subscription.unsubscribe()
+                        return
+
+                    # First miss: the batch may simply have been lost, so reset
+                    # the pending counts and ask again before giving up.
                     self._pending_messages = 0
                     self._pending_bytes = 0
                     await self._send_request()

nats-jetstream/tests/test_consumer.py

@@ -1084,6 +1084,34 @@ async def test_consumer_info_timestamp(jetstream: JetStream):
     assert info.timestamp.tzinfo is not None
 
 
+@pytest.mark.asyncio
+async def test_messages_end_after_repeated_missed_heartbeats(jetstream: JetStream):
+    """A stream whose consumer stops answering ends instead of waiting forever.
+
+    One missed heartbeat window can be a lost batch, so the stream asks again.
+    A second means nothing is answering, and the iterator has to end -- an
+    ordered consumer takes that as its cue to reset.
+    """
+    stream = await jetstream.create_stream(name="hb_miss", subjects=["HBM.*"])
+    consumer = await stream.create_consumer(name="hb_miss", durable_name="hb_miss", ack_policy="explicit")
+
+    messages = await consumer.messages(heartbeat=0.2, max_wait=5.0)
+    pending = asyncio.create_task(messages.__anext__())
+    await asyncio.sleep(0.1)
+
+    # Simulate the consumer going silent while the connection stays up: stop the
+    # request loop and neuter the retry, so both heartbeat windows lapse unanswered.
+    async def _silent() -> None:
+        return None
+
+    if messages._request_task is not None:
+        messages._request_task.cancel()
+    messages._send_request = _silent  # type: ignore[method-assign]
+
+    with pytest.raises(StopAsyncIteration):
+        await asyncio.wait_for(pending, timeout=10)
+
+
 @pytest.mark.asyncio
 async def test_consumer_reset_to_seq(jetstream: JetStream):
     """Reset a consumer forward to a specific stream sequence (ADR-60).

nats-jetstream/tests/test_ordered_consumer.py

@@ -757,6 +757,47 @@ async def test_ordered_consumer_close_connection(server: Server):
     assert len(received) == 3
 
 
+@pytest.mark.parametrize(
+    ("max_wait", "expected_heartbeat"),
+    [(None, 15.0), (10.0, 5.0), (4.0, 2.0)],
+)
+async def test_ordered_consumer_messages_request_heartbeats(
+    jetstream: JetStream, max_wait: float | None, expected_heartbeat: float
+):
+    """messages() always asks for idle heartbeats, kept under the request expiry.
+
+    An ordered consumer only resets when its inner stream ends. Heartbeats are
+    what end it when the consumer stops answering, so leaving them unset means
+    the iterator waits forever.
+    """
+    stream = await jetstream.create_stream(name="oc_heartbeat", subjects=["OCHB.*"])
+    consumer = await stream.ordered_consumer(filter_subjects=["OCHB.*"])
+
+    messages = await consumer.messages(max_wait=max_wait)
+    inner = await messages._create_inner_stream()
+    try:
+        assert inner._heartbeat == expected_heartbeat
+        # The server rejects a request whose heartbeat is not under its expiry.
+        assert inner._heartbeat < inner._expires
+    finally:
+        await inner.stop()
+        await messages.stop()
+
+
+async def test_ordered_consumer_messages_honours_explicit_heartbeat(jetstream: JetStream):
+    """An explicit heartbeat is passed through untouched."""
+    stream = await jetstream.create_stream(name="oc_heartbeat_explicit", subjects=["OCHBE.*"])
+    consumer = await stream.ordered_consumer(filter_subjects=["OCHBE.*"])
+
+    messages = await consumer.messages(max_wait=20.0, heartbeat=1.5)
+    inner = await messages._create_inner_stream()
+    try:
+        assert inner._heartbeat == 1.5
+    finally:
+        await inner.stop()
+        await messages.stop()
+
+
 async def test_ordered_consumer_messages_server_restart(server: Server, store_dir: str):
     """messages() recovers after a server restart."""
     client = await connect(server.client_url, reconnect_max_attempts=0)
@@ -796,11 +837,16 @@ async def test_ordered_consumer_messages_server_restart(server: Server, store_di
         for i in range(5):
             await js.publish(f"OC.SRV.more{i}", f"more {i}".encode())
 
-        # Should recover and deliver new messages
-        async for msg in messages:
-            received.append(msg.data.decode())
-            if len(received) == 10:
-                break
+        # Should recover and deliver new messages. Bounded: if recovery never
+        # happens this iterator has nothing to end it, and an unbounded wait
+        # here stalls the whole run rather than failing.
+        async def drain_remaining() -> None:
+            async for msg in messages:
+                received.append(msg.data.decode())
+                if len(received) == 10:
+                    break
+
+        await asyncio.wait_for(drain_remaining(), timeout=60)
 
         assert len(received) == 10
         await messages.stop()