diff --git a/nats/src/nats/js/api.py b/nats/src/nats/js/api.py index e559e7b4..771332f0 100644 --- a/nats/src/nats/js/api.py +++ b/nats/src/nats/js/api.py @@ -37,6 +37,7 @@ class Header(str, Enum): LAST_CONSUMER = "Nats-Last-Consumer" LAST_STREAM = "Nats-Last-Stream" MSG_ID = "Nats-Msg-Id" + PIN_ID = "Nats-Pin-Id" MSG_TTL = "Nats-TTL" ROLLUP = "Nats-Rollup" SCHEDULE = "Nats-Schedule" @@ -67,6 +68,7 @@ class StatusCode(str, Enum): NO_MESSAGES = "404" REQUEST_TIMEOUT = "408" CONFLICT = "409" + PIN_ID_MISMATCH = "423" CONTROL_MESSAGE = "100" @@ -620,6 +622,32 @@ class ReplayPolicy(str, Enum): ORIGINAL = "original" +class PriorityPolicy(str, Enum): + """Priority policy for pull consumer priority groups. + + Enables flexible failover and priority management when multiple clients + are pulling from the same consumer. + + Introduced in nats-server 2.11.0 (``PRIORITIZED`` in 2.12.0). + + References: + * `Consumers, Pull consumer priority groups ` + * `Consumers, Prioritized pull consumer policy ` + """ # noqa: E501 + + NONE = "none" + """Default, no priority handling.""" + PINNED = "pinned_client" + """Pins the consumer to a single client per group; others take over when it goes away.""" + OVERFLOW = "overflow" + """Only delivers to a client once ``min_pending`` or ``min_ack_pending`` thresholds are reached.""" + PRIORITIZED = "prioritized" + """Delivers to the client with the highest priority (0-9, 0 is highest and default). + + Introduced in nats-server 2.12.0. + """ + + @dataclass class ConsumerConfig(Base): """Consumer configuration. @@ -672,12 +700,26 @@ class ConsumerConfig(Base): # Introduced in nats-server 2.11.0. pause_until: Optional[str] = None + # Priority policy. + # Introduced in nats-server 2.11.0. + priority_policy: Optional[PriorityPolicy] = None + + # The duration (seconds) after which the client will be unpinned if no new + # pull requests are sent. Used with PriorityPolicy.PINNED. + # Introduced in nats-server 2.11.0. + priority_timeout: Optional[float] = None + + # Priority groups this consumer supports. + # Introduced in nats-server 2.11.0. + priority_groups: Optional[list[str]] = None + @classmethod def from_response(cls, resp: Dict[str, Any]): cls._convert_nanoseconds(resp, "ack_wait") cls._convert_nanoseconds(resp, "idle_heartbeat") cls._convert_nanoseconds(resp, "inactive_threshold") cls._convert_utc_iso(resp, "opt_start_time") + cls._convert_nanoseconds(resp, "priority_timeout") if "backoff" in resp: resp["backoff"] = [val / _NANOSECOND for val in resp["backoff"]] return super().from_response(resp) @@ -689,6 +731,8 @@ def as_dict(self) -> Dict[str, object]: result["ack_wait"] = self._to_nanoseconds(self.ack_wait) result["idle_heartbeat"] = self._to_nanoseconds(self.idle_heartbeat) result["inactive_threshold"] = self._to_nanoseconds(self.inactive_threshold) + if self.priority_timeout is not None: + result["priority_timeout"] = self._to_nanoseconds(self.priority_timeout) if self.backoff: result["backoff"] = [self._to_nanoseconds(i) for i in self.backoff] return result @@ -712,6 +756,32 @@ def as_dict(self) -> Dict[str, object]: return result +@dataclass +class PriorityGroupState(Base): + """ + State of a consumer priority group. + + Introduced in nats-server 2.11.0. + """ + + group: str + # Generated ID of the pinned client. Only set when a client is pinned. + pinned_client_id: Optional[str] = None + # When the client was pinned. Only set when a client is pinned. + pinned_ts: Optional[datetime.datetime] = None + + @classmethod + def from_response(cls, resp: Dict[str, Any]): + cls._convert_utc_iso(resp, "pinned_ts") + return super().from_response(resp) + + def as_dict(self) -> Dict[str, object]: + result = super().as_dict() + if self.pinned_ts is not None: + result["pinned_ts"] = self._to_utc_iso(self.pinned_ts) + return result + + @dataclass class ConsumerInfo(Base): """ @@ -736,6 +806,8 @@ class ConsumerInfo(Base): # RFC 3339 timestamp until which the consumer is paused. # Introduced in nats-server 2.11.0. pause_remaining: Optional[str] = None + # Introduced in nats-server 2.11.0. + priority_groups: Optional[list[PriorityGroupState]] = None @classmethod def from_response(cls, resp: Dict[str, Any]): @@ -743,6 +815,7 @@ def from_response(cls, resp: Dict[str, Any]): cls._convert(resp, "ack_floor", SequenceInfo) cls._convert(resp, "config", ConsumerConfig) cls._convert(resp, "cluster", ClusterInfo) + cls._convert(resp, "priority_groups", PriorityGroupState) cls._convert_utc_iso(resp, "created") return super().from_response(resp) diff --git a/nats/src/nats/js/client.py b/nats/src/nats/js/client.py index d516a6e2..d63118a4 100644 --- a/nats/src/nats/js/client.py +++ b/nats/src/nats/js/client.py @@ -546,6 +546,7 @@ async def pull_subscribe( pending_msgs_limit: int = DEFAULT_JS_SUB_PENDING_MSGS_LIMIT, pending_bytes_limit: int = DEFAULT_JS_SUB_PENDING_BYTES_LIMIT, inbox_prefix: Optional[bytes] = None, + priority_group: Optional[str] = None, ) -> JetStreamContext.PullSubscription: """Create consumer and pull subscription. @@ -580,6 +581,9 @@ async def main(): if stream is None: stream = await self._jsm.find_stream_name_by_subject(subject) + if config and config.priority_groups and priority_group is None: + raise ValueError("nats: priority_group is required when consumer has priority_groups configured") + should_create = True try: if durable: @@ -605,6 +609,10 @@ async def main(): consumer_name = self._nc._nuid.next().decode() config.name = consumer_name + # Auto created consumers use the priority group, unless priority_groups is set. + if not config.priority_groups and priority_group: + config.priority_groups = [priority_group] + await self._jsm.add_consumer(stream, config=config) return await self.pull_subscribe_bind( @@ -614,6 +622,7 @@ async def main(): pending_bytes_limit=pending_bytes_limit, pending_msgs_limit=pending_msgs_limit, name=consumer_name, + priority_group=priority_group, ) async def pull_subscribe_bind( @@ -625,6 +634,7 @@ async def pull_subscribe_bind( pending_bytes_limit: int = DEFAULT_JS_SUB_PENDING_BYTES_LIMIT, name: Optional[str] = None, durable: Optional[str] = None, + priority_group: Optional[str] = None, ) -> JetStreamContext.PullSubscription: """ pull_subscribe returns a `PullSubscription` that can be delivered messages @@ -680,6 +690,7 @@ async def main(): stream=stream, consumer=consumer_name, deliver=deliver, + group=priority_group, ) @classmethod @@ -703,11 +714,16 @@ def _is_temporary_error(cls, status: Optional[str]) -> bool: status == api.StatusCode.NO_MESSAGES or status == api.StatusCode.CONFLICT or status == api.StatusCode.REQUEST_TIMEOUT + or status == api.StatusCode.PIN_ID_MISMATCH ): return True else: return False + @classmethod + def _is_pin_id_mismatch_error(cls, status: Optional[str]) -> bool: + return status == api.StatusCode.PIN_ID_MISMATCH + @classmethod def _is_heartbeat(cls, status: Optional[str]) -> bool: if status == api.StatusCode.CONTROL_MESSAGE: @@ -997,6 +1013,7 @@ def __init__( stream: str, consumer: str, deliver: bytes, + group: Optional[str] = None, ) -> None: # JS/JSM context self._js = js @@ -1009,6 +1026,8 @@ def __init__( prefix = self._js._prefix self._nms = f"{prefix}.CONSUMER.MSG.NEXT.{stream}.{consumer}" self._deliver = deliver.decode() + self._pin_id: Optional[str] = None + self._group = group @property def pending_msgs(self) -> int: @@ -1055,6 +1074,9 @@ async def fetch( batch: int = 1, timeout: Optional[float] = 5, heartbeat: Optional[float] = None, + min_pending: Optional[int] = None, + min_ack_pending: Optional[int] = None, + priority: Optional[int] = None, ) -> List[Msg]: """ fetch makes a request to JetStream to be delivered a set of messages. @@ -1062,6 +1084,12 @@ async def fetch( :param batch: Number of messages to fetch from server. :param timeout: Max duration of the fetch request before it expires. :param heartbeat: Idle Heartbeat interval in seconds for the fetch request. + :param min_pending: Only deliver when the consumer has at least this many + pending messages. Requires ``PriorityPolicy.OVERFLOW``. + :param min_ack_pending: Only deliver when the consumer has at least this + many unacknowledged messages. Requires ``PriorityPolicy.OVERFLOW``. + :param priority: Priority of this request from 0 (highest) to 9. + Requires ``PriorityPolicy.PRIORITIZED``. :: @@ -1092,12 +1120,18 @@ async def main(): raise ValueError("nats: invalid batch size") if timeout is not None and timeout <= 0: raise ValueError("nats: invalid fetch timeout") + if min_pending is not None and min_pending <= 0: + raise ValueError("nats: min_pending must be more than 0") + if min_ack_pending is not None and min_ack_pending <= 0: + raise ValueError("nats: min_ack_pending must be more than 0") + if priority is not None and not (0 <= priority <= 9): + raise ValueError("nats: priority must be 0-9") expires = int(timeout * 1_000_000_000) - 100_000 if timeout else None if batch == 1: - msg = await self._fetch_one(expires, timeout, heartbeat) + msg = await self._fetch_one(expires, timeout, heartbeat, min_pending, min_ack_pending, priority) return [msg] - msgs = await self._fetch_n(batch, expires, timeout, heartbeat) + msgs = await self._fetch_n(batch, expires, timeout, heartbeat, min_pending, min_ack_pending, priority) return msgs async def _fetch_one( @@ -1105,6 +1139,9 @@ async def _fetch_one( expires: Optional[int], timeout: Optional[float], heartbeat: Optional[float] = None, + min_pending: Optional[int] = None, + min_ack_pending: Optional[int] = None, + priority: Optional[int] = None, ) -> Msg: queue = self._sub._pending_queue @@ -1124,21 +1161,35 @@ async def _fetch_one( pass # Make lingering request with expiration and wait for response. - next_req = {} - next_req["batch"] = 1 - if expires: - next_req["expires"] = int(expires) - if heartbeat: - next_req["idle_heartbeat"] = int(heartbeat * 1_000_000_000) # to nanoseconds - - await self._nc.publish( - self._nms, - json.dumps(next_req).encode(), - self._deliver, - ) + async def send_next_request() -> None: + next_req = {} + next_req["batch"] = 1 + if expires: + next_req["expires"] = int(expires) + if heartbeat: + next_req["idle_heartbeat"] = int(heartbeat * 1_000_000_000) # to nanoseconds + if self._group: + next_req["group"] = self._group + pin_id = self.pin_id + if pin_id: + next_req["id"] = pin_id + if min_pending is not None: + next_req["min_pending"] = min_pending + if min_ack_pending is not None: + next_req["min_ack_pending"] = min_ack_pending + if priority is not None: + next_req["priority"] = priority + await self._nc.publish( + self._nms, + json.dumps(next_req).encode(), + self._deliver, + ) + + await send_next_request() start_time = time.monotonic() got_any_response = False + resent_without_pin_id = False while True: try: deadline = JetStreamContext._time_until(timeout, start_time) @@ -1152,6 +1203,19 @@ async def _fetch_one( got_any_response = True continue + if JetStreamContext._is_pin_id_mismatch_error(status): + # The pin this request carried is no longer valid and + # the server has discarded the request. Drop the stale + # id and re-issue once without it so the fetch can + # recover within its own deadline instead of waiting + # on a request that will never be served. + self.pin_id = None + got_any_response = True + if not resent_without_pin_id: + resent_without_pin_id = True + await send_next_request() + continue + # In case of a temporary error, treat it as a timeout to retry. if JetStreamContext._is_temporary_error(status): raise nats.errors.TimeoutError @@ -1159,6 +1223,9 @@ async def _fetch_one( # Any other type of status message is an error. raise nats.js.errors.APIError.from_msg(msg) else: + pin_id = msg.headers.get(api.Header.PIN_ID) if msg.headers else None + if pin_id: + self.pin_id = pin_id return msg except asyncio.TimeoutError: deadline = JetStreamContext._time_until(timeout, start_time) @@ -1177,6 +1244,9 @@ async def _fetch_n( expires: Optional[int], timeout: Optional[float], heartbeat: Optional[float] = None, + min_pending: Optional[int] = None, + min_ack_pending: Optional[int] = None, + priority: Optional[int] = None, ) -> List[Msg]: msgs = [] queue = self._sub._pending_queue @@ -1218,6 +1288,17 @@ async def _fetch_n( if heartbeat: next_req["idle_heartbeat"] = int(heartbeat * 1_000_000_000) # to nanoseconds next_req["no_wait"] = True + if self._group: + next_req["group"] = self._group + pin_id = self.pin_id + if pin_id: + next_req["id"] = pin_id + if min_pending is not None: + next_req["min_pending"] = min_pending + if min_ack_pending is not None: + next_req["min_ack_pending"] = min_ack_pending + if priority is not None: + next_req["priority"] = priority await self._nc.publish( self._nms, json.dumps(next_req).encode(), @@ -1241,8 +1322,13 @@ async def _fetch_n( # a possible i/o timeout error or due to a disconnection. got_any_response = True pass + elif JetStreamContext._is_pin_id_mismatch_error(status): + self.pin_id = None elif JetStreamContext._is_processable_msg(status, msg): # First processable message received, do not raise error from now. + pin_id = msg.headers.get(api.Header.PIN_ID) if msg.headers else None + if pin_id: + self.pin_id = pin_id msgs.append(msg) needed -= 1 @@ -1259,7 +1345,12 @@ async def _fetch_n( # Skip heartbeats. got_any_response = True continue + elif JetStreamContext._is_pin_id_mismatch_error(status): + self.pin_id = None elif JetStreamContext._is_processable_msg(status, msg): + pin_id = msg.headers.get(api.Header.PIN_ID) if msg.headers else None + if pin_id: + self.pin_id = pin_id needed -= 1 msgs.append(msg) except asyncio.TimeoutError: @@ -1299,7 +1390,17 @@ async def _fetch_n( next_req["expires"] = expires if heartbeat: next_req["idle_heartbeat"] = int(heartbeat * 1_000_000_000) # to nanoseconds - + if self._group: + next_req["group"] = self._group + pin_id = self.pin_id + if pin_id: + next_req["id"] = pin_id + if min_pending is not None: + next_req["min_pending"] = min_pending + if min_ack_pending is not None: + next_req["min_ack_pending"] = min_ack_pending + if priority is not None: + next_req["priority"] = priority await self._nc.publish( self._nms, json.dumps(next_req).encode(), @@ -1338,8 +1439,13 @@ async def _fetch_n( if JetStreamContext._is_heartbeat(status): got_any_response = True continue + if JetStreamContext._is_pin_id_mismatch_error(status): + self.pin_id = None if not status: + pin_id = msg.headers.get(api.Header.PIN_ID) if msg.headers else None + if pin_id: + self.pin_id = pin_id needed -= 1 msgs.append(msg) break @@ -1363,6 +1469,8 @@ async def _fetch_n( if JetStreamContext._is_heartbeat(status): got_any_response = True continue + if JetStreamContext._is_pin_id_mismatch_error(status): + self.pin_id = None if status in ( api.StatusCode.NO_MESSAGES, api.StatusCode.REQUEST_TIMEOUT, @@ -1371,6 +1479,9 @@ async def _fetch_n( # request; return what we have. break if JetStreamContext._is_processable_msg(status, msg): + pin_id = msg.headers.get(api.Header.PIN_ID) if msg.headers else None + if pin_id: + self.pin_id = pin_id needed -= 1 msgs.append(msg) except asyncio.TimeoutError: @@ -1383,6 +1494,14 @@ async def _fetch_n( return msgs + @property + def pin_id(self) -> Optional[str]: + return self._pin_id + + @pin_id.setter + def pin_id(self, pin_id: Optional[str]) -> None: + self._pin_id = pin_id + ###################### # # # KeyValue Context # diff --git a/nats/src/nats/js/errors.py b/nats/src/nats/js/errors.py index 722a713f..c2ee694b 100644 --- a/nats/src/nats/js/errors.py +++ b/nats/src/nats/js/errors.py @@ -83,6 +83,8 @@ def from_error(cls, err: Dict[str, Any]): raise ServiceUnavailableError(**err) elif code == 500: raise ServerError(**err) + elif code == 423: + raise PinIdMismatchError(**err) elif code == 404: raise NotFoundError(**err) elif code == 400: @@ -112,6 +114,17 @@ class ServerError(APIError): pass +class PinIdMismatchError(APIError): + """ + A 423 error + + PinIdMismatchError is returned when Pin ID sent in the request does not match + the currently pinned consumer subscriber ID on the server. + """ + + pass + + class NotFoundError(APIError): """ A 404 error diff --git a/nats/src/nats/js/manager.py b/nats/src/nats/js/manager.py index 7272ad3e..267cea97 100644 --- a/nats/src/nats/js/manager.py +++ b/nats/src/nats/js/manager.py @@ -523,6 +523,21 @@ async def get_last_msg( """ return await self.get_msg(stream_name, subject=subject, direct=direct) + async def unpin_consumer(self, stream_name: str, consumer_name: str, group: str) -> None: + """ + unpin_consumer releases the client currently pinned to a priority group + of a consumer using ``PriorityPolicy.PINNED``, so that the next pull + request from any client in that group is pinned instead. + + :param stream_name: Name of the stream the consumer belongs to. + :param consumer_name: Name of the consumer. + :param group: Priority group to unpin. + """ + req_subject = f"{self._prefix}.CONSUMER.UNPIN.{stream_name}.{consumer_name}" + req = {"group": group} + data = json.dumps(req) + _ = await self._api_request(req_subject, data.encode()) + async def _api_request( self, req_subject: str, diff --git a/nats/tests/test_js.py b/nats/tests/test_js.py index eb7610f6..a3f4fa4d 100644 --- a/nats/tests/test_js.py +++ b/nats/tests/test_js.py @@ -6241,3 +6241,465 @@ async def error_handler(e): assert e.description == "changed" await nc.close() + + +class PriorityGroupsApiTest(unittest.TestCase): + """Unit tests for ADR-42 priority group types.""" + + def test_consumer_config_round_trip(self): + config = nats.js.api.ConsumerConfig( + priority_policy=nats.js.api.PriorityPolicy.PINNED, + priority_groups=["A", "B"], + priority_timeout=30, + ) + d = config.as_dict() + assert d["priority_policy"] == "pinned_client" + assert d["priority_groups"] == ["A", "B"] + assert d["priority_timeout"] == 30_000_000_000 + parsed = nats.js.api.ConsumerConfig.from_response(json.loads(json.dumps(d))) + assert parsed.priority_policy == nats.js.api.PriorityPolicy.PINNED + assert parsed.priority_groups == ["A", "B"] + assert parsed.priority_timeout == 30 + + def test_consumer_config_omits_unset_priority_fields(self): + d = nats.js.api.ConsumerConfig().as_dict() + assert "priority_policy" not in d + assert "priority_groups" not in d + assert "priority_timeout" not in d + + def test_consumer_info_priority_groups_unpinned(self): + # The server omits pinned_client_id and pinned_ts until a client is pinned. + info = nats.js.api.ConsumerInfo.from_response( + { + "stream_name": "S", + "name": "C", + "created": "2026-08-29T09:00:00Z", + "config": {"priority_policy": "pinned_client", "priority_groups": ["A"]}, + "delivered": {"consumer_seq": 0, "stream_seq": 0}, + "ack_floor": {"consumer_seq": 0, "stream_seq": 0}, + "num_ack_pending": 0, + "num_redelivered": 0, + "num_waiting": 0, + "num_pending": 0, + "priority_groups": [{"group": "A"}], + } + ) + assert info.priority_groups == [nats.js.api.PriorityGroupState(group="A")] + + def test_consumer_info_priority_groups_pinned(self): + info = nats.js.api.ConsumerInfo.from_response( + { + "stream_name": "S", + "name": "C", + "created": "2026-08-29T09:00:00Z", + "config": {"priority_policy": "pinned_client", "priority_groups": ["A"]}, + "delivered": {"consumer_seq": 0, "stream_seq": 0}, + "ack_floor": {"consumer_seq": 0, "stream_seq": 0}, + "num_ack_pending": 0, + "num_redelivered": 0, + "num_waiting": 0, + "num_pending": 0, + "priority_groups": [ + {"group": "A", "pinned_client_id": "abc", "pinned_ts": "2026-08-29T10:00:00Z"}, + ], + } + ) + state = info.priority_groups[0] + assert state.group == "A" + assert state.pinned_client_id == "abc" + assert state.pinned_ts == datetime.datetime(2026, 8, 29, 10, 0, tzinfo=datetime.timezone.utc) + + +class PriorityGroupsFeaturesTest(SingleJetStreamServerTestCase): + @async_test + async def test_consumer_overflow(self): + nc = await nats.connect() + + server_version = nc.connected_server_version + if server_version.major == 2 and server_version.minor < 11: + pytest.skip("consumer group overflow requires nats-server v2.11.0 or later") + + js = nc.jetstream() + + # create stream + await js.add_stream( + name="PRIORITIES", + subjects=["foo"], + ) + + # create consumer with overflow priority policy + cinfo = await js.add_consumer( + "PRIORITIES", + nats.js.api.ConsumerConfig( + priority_policy=nats.js.api.PriorityPolicy.OVERFLOW, + priority_groups=["A"], + ), + ) + assert cinfo.config.priority_policy == nats.js.api.PriorityPolicy.OVERFLOW + + # 1. Below threshold - no messages delivered + # - publish 100 msgs + # - fetch 10 msgs with min_pending 110 + # - should not get any msgs since 100<110 + psub = await js.pull_subscribe_bind( + cinfo.name, + cinfo.stream_name, + priority_group="A", + ) + for i in range(0, 100): + await js.publish("foo", f"{i}".encode()) + with pytest.raises(TimeoutError): + await psub.fetch(10, timeout=0.5, min_pending=110) + await psub.unsubscribe() + + # 2. Above threshold - messages delivered + # - publish 100 more msgs + # - fetch 10 msgs with min_pending 110 + # - should get 10 msgs since (200-10)>110 + psub = await js.pull_subscribe_bind( + cinfo.name, + cinfo.stream_name, + priority_group="A", + ) + for i in range(0, 100): + await js.publish("foo", f"{i}".encode()) + msgs = await psub.fetch(10, timeout=0.5, min_pending=110) + assert len(msgs) == 10 + for msg in msgs: # clean up + await msg.ack_sync() + await psub.unsubscribe() + + # 3: MinAckPending - no unacked messages yet + # - fetch 10 msgs with min_ack_pending 10 + # - should get 0 msgs since no pending acks currently + psub = await js.pull_subscribe_bind( + cinfo.name, + cinfo.stream_name, + priority_group="A", + ) + with pytest.raises(TimeoutError): + await psub.fetch(10, timeout=0.5, min_ack_pending=10) + await psub.unsubscribe() + + # 4: MinAckPending threshold met + # - create 10 pending acks + # - fetch 10 msgs with min_ack_pending 10 + # - should get 10 msgs since 10 pending acks >=10 + # NOTE: the psub's buffer queue can get filled with extra messages which + # leak into subsequent fetch calls, so to check unbuffered behavior we + # use separate subs + psub1 = await js.pull_subscribe_bind( + cinfo.name, + cinfo.stream_name, + priority_group="A", + ) + psub2 = await js.pull_subscribe_bind( + cinfo.name, + cinfo.stream_name, + priority_group="A", + ) + unacked_msgs = await psub1.fetch(10, timeout=0.5) + msgs = await psub2.fetch(10, timeout=0.5, min_ack_pending=10) + assert len(msgs) == 10 + for msg in unacked_msgs + msgs: # clean up + await msg.ack_sync() + + await nc.close() + + @async_test + async def test_consumer_pinned(self): + nc = await nats.connect() + + server_version = nc.connected_server_version + if server_version.major == 2 and server_version.minor < 11: + pytest.skip("consumer group pinning requires nats-server v2.11.0 or later") + + js = nc.jetstream() + + # create stream + await js.add_stream( + name="PRIORITIES", + subjects=["foo"], + ) + + # create consumer with pinned priority policy + cinfo = await js.add_consumer( + "PRIORITIES", + nats.js.api.ConsumerConfig( + priority_policy=nats.js.api.PriorityPolicy.PINNED, + priority_timeout=1.0, + priority_groups=["A"], + ), + ) + assert cinfo.config.priority_policy == nats.js.api.PriorityPolicy.PINNED + + # publish messages + for i in range(100): + await js.publish("foo", f"{i}".encode()) + + # 1. Priority group validation - invalid group + psub = await js.pull_subscribe_bind( + cinfo.name, + cinfo.stream_name, + priority_group="BAD", + ) + with pytest.raises(nats.js.errors.APIError, match="Invalid Priority Group"): + await psub.fetch(10, timeout=0.5) + await psub.unsubscribe() + + # 2. Priority group validation - no group + psub = await js.pull_subscribe_bind( + cinfo.name, + cinfo.stream_name, + ) + with pytest.raises(nats.js.errors.APIError, match="Priority Group missing"): + await psub.fetch(10, timeout=0.5) + await psub.unsubscribe() + + # 3. First consumer gets pinned + psub1 = await js.pull_subscribe_bind( + cinfo.name, + cinfo.stream_name, + priority_group="A", + ) + msgs = await psub1.fetch(10, timeout=0.5) + assert len(msgs) == 10 + first_pin_id = msgs[0].headers.get("Nats-Pin-Id") if msgs[0].headers else None + assert first_pin_id is not None + # all messages should have same pin id + for msg in msgs: + assert msg.headers.get("Nats-Pin-Id") == first_pin_id + await msg.ack_sync() + + # 4. Different consumer instance can't fetch while pinned + psub2 = await js.pull_subscribe_bind( + cinfo.name, + cinfo.stream_name, + priority_group="A", + ) + with pytest.raises(TimeoutError): + await psub2.fetch(10, timeout=0.5) + + # 5. Original consumer continues to work + msgs = await psub1.fetch(10, timeout=0.5) + assert len(msgs) == 10 + for msg in msgs: + assert msg.headers.get("Nats-Pin-Id") == first_pin_id + await msg.ack_sync() + + # 6. After TTL expires, pin ID changes + await asyncio.sleep(1.5) # longer than priority_timeout (1s) + msgs = await psub1.fetch(10, timeout=0.5) + assert len(msgs) == 10 + new_pin_id = msgs[0].headers.get("Nats-Pin-Id") if msgs[0].headers else None + assert new_pin_id is not None + assert new_pin_id != first_pin_id + for msg in msgs: # clean up + await msg.ack_sync() + + await psub1.unsubscribe() + await psub2.unsubscribe() + await nc.close() + + @async_test + async def test_consumer_unpin(self): + nc = await nats.connect() + + server_version = nc.connected_server_version + if server_version.major == 2 and server_version.minor < 11: + pytest.skip("consumer group unpinning requires nats-server v2.11.0 or later") + + js = nc.jetstream() + jsm = js._jsm + + # create stream + await js.add_stream( + name="PRIORITIES", + subjects=["foo"], + ) + + # create consumer with pinned priority policy and long TTL + cinfo = await js.add_consumer( + "PRIORITIES", + nats.js.api.ConsumerConfig( + priority_policy=nats.js.api.PriorityPolicy.PINNED, + priority_timeout=50.0, + priority_groups=["A"], + ), + ) + assert cinfo.config.priority_policy == nats.js.api.PriorityPolicy.PINNED + + # publish messages + for i in range(100): + await js.publish("foo", f"{i}".encode()) + + # 1. First consumer gets pinned + psub1 = await js.pull_subscribe_bind( + cinfo.name, + cinfo.stream_name, + priority_group="A", + ) + msgs = await psub1.fetch(1, timeout=0.5) + assert len(msgs) == 1 + first_pin_id = msgs[0].headers.get("Nats-Pin-Id") if msgs[0].headers else None + assert first_pin_id is not None + await msgs[0].ack_sync() + + # 2. Second consumer can't get messages while first is pinned + psub2 = await js.pull_subscribe_bind( + cinfo.name, + cinfo.stream_name, + priority_group="A", + ) + with pytest.raises(TimeoutError): + await psub2.fetch(1, timeout=0.5) + await psub2.unsubscribe() + + # 3. Manual unpin allows third consumer + psub3 = await js.pull_subscribe_bind( + cinfo.name, + cinfo.stream_name, + priority_group="A", + ) + + # Unpin the consumer + await jsm.unpin_consumer(cinfo.stream_name, cinfo.name, "A") + + # Third consumer should now receive message with new pin ID + msgs = await psub3.fetch(1, timeout=0.5) + assert len(msgs) == 1 + new_pin_id = msgs[0].headers.get("Nats-Pin-Id") if msgs[0].headers else None + assert new_pin_id is not None + assert new_pin_id != first_pin_id + + await psub1.unsubscribe() + await psub3.unsubscribe() + + # 4. Test unpin on non-existent consumer + with pytest.raises(nats.js.errors.NotFoundError): + await jsm.unpin_consumer("PRIORITIES", "nonexistent", "A") + + await nc.close() + + @async_test + async def test_consumer_pin_id_mismatch_recovers(self): + """A stale pin id is dropped and the fetch recovers within its own deadline.""" + nc = await nats.connect() + + server_version = nc.connected_server_version + if server_version.major == 2 and server_version.minor < 11: + pytest.skip("consumer group pinning requires nats-server v2.11.0 or later") + + js = nc.jetstream() + jsm = js._jsm + + await js.add_stream(name="PRIORITIES", subjects=["foo"]) + cinfo = await js.add_consumer( + "PRIORITIES", + nats.js.api.ConsumerConfig( + priority_policy=nats.js.api.PriorityPolicy.PINNED, + priority_timeout=50.0, + priority_groups=["A"], + ), + ) + for i in range(10): + await js.publish("foo", f"{i}".encode()) + + # Both the single-message and batched paths must recover from a 423. + for batch in (1, 2): + psub = await js.pull_subscribe_bind( + cinfo.name, + cinfo.stream_name, + priority_group="A", + ) + msgs = await psub.fetch(1, timeout=1.0) + for msg in msgs: + await msg.ack_sync() + stale_pin_id = psub.pin_id + assert stale_pin_id is not None + + # Unpinning invalidates the id this subscription still holds, so the + # next pull request is rejected with a 423. + await jsm.unpin_consumer(cinfo.stream_name, cinfo.name, "A") + assert psub.pin_id == stale_pin_id + + msgs = await psub.fetch(batch, timeout=2.0) + assert len(msgs) == batch + for msg in msgs: + await msg.ack_sync() + assert psub.pin_id is not None + assert psub.pin_id != stale_pin_id + + await psub.unsubscribe() + await jsm.unpin_consumer(cinfo.stream_name, cinfo.name, "A") + + await nc.close() + + @async_test + async def test_consumer_prioritized(self): + nc = await nats.connect() + + server_version = nc.connected_server_version + if server_version.major == 2 and server_version.minor < 12: + pytest.skip("consumer group priority requires nats-server v2.12.0 or later") + + js = nc.jetstream() + + # create stream + await js.add_stream( + name="PRIORITIES", + subjects=["foo"], + ) + + # create consumer with prioritized priority policy + cinfo = await js.add_consumer( + "PRIORITIES", + nats.js.api.ConsumerConfig( + priority_policy=nats.js.api.PriorityPolicy.PRIORITIZED, + priority_groups=["A"], + ), + ) + assert cinfo.config.priority_policy == nats.js.api.PriorityPolicy.PRIORITIZED + + # Test: Messages distributed based on priority + # Higher priority (lower number) consumers get messages first + + # Create two consumer instances + psub1 = await js.pull_subscribe_bind( + cinfo.name, + cinfo.stream_name, + priority_group="A", + ) + psub2 = await js.pull_subscribe_bind( + cinfo.name, + cinfo.stream_name, + priority_group="A", + ) + + # publish 100 messages + for i in range(100): + await js.publish("foo", f"{i}".encode()) + + # Start concurrent fetches: + # psub1 with priority=1 (lower priority) requesting 100 messages + # psub2 with priority=0 (higher priority) requesting 75 messages + # Expected: psub2 gets 75 first, psub1 gets remaining 25 + + fetch1_task = asyncio.create_task(psub1.fetch(100, timeout=2.0, priority=1)) + fetch2_task = asyncio.create_task(psub2.fetch(75, timeout=2.0, priority=0)) + + # Wait for both fetches + msgs1, msgs2 = await asyncio.gather(fetch1_task, fetch2_task) + + # psub2 (priority 0) should get 75 messages + assert len(msgs2) == 75 + + # psub1 (priority 1) should get remaining 25 messages + assert len(msgs1) == 25 + + for msg in msgs1 + msgs2: # clean up + await msg.ack_sync() + + await psub1.unsubscribe() + await psub2.unsubscribe() + await nc.close()