Skip to content

Commit 1e87e76

Browse files
aKardaszLancetnik
andauthored
Resolve Issue 1386, Add rpc_prefix (#1484)
* update adding nuid usage as well as reply_to_prefix * add missing commas * update reply_to nuid generation, add missing imports * rename reply_to_prefix to rpc_prefix * refactor rpc reply_to address generation. * remove reply_to, update nats implentation * add back reply_to * remove rpc_prefix, update NatsJSFastProducer * Add nats client inbox_prefix test * fixes to inbox_prefix test * lint: fix bandit --------- Co-authored-by: Pastukhov Nikita <nikita@pastukhov-dev.ru> Co-authored-by: Nikita Pastukhov <diementros@yandex.ru>
1 parent dfec397 commit 1e87e76

File tree

6 files changed

+99
-10
lines changed

6 files changed

+99
-10
lines changed

docs/docs/SUMMARY.md

+2
Original file line numberDiff line numberDiff line change
@@ -1007,6 +1007,8 @@ search:
10071007
- [to_async](api/faststream/utils/functions/to_async.md)
10081008
- no_cast
10091009
- [NoCast](api/faststream/utils/no_cast/NoCast.md)
1010+
- nuid
1011+
- [NUID](api/faststream/utils/nuid/NUID.md)
10101012
- path
10111013
- [compile_path](api/faststream/utils/path/compile_path.md)
10121014
- Contributing
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
---
2+
# 0.5 - API
3+
# 2 - Release
4+
# 3 - Contributing
5+
# 5 - Template Page
6+
# 10 - Default
7+
search:
8+
boost: 0.5
9+
---
10+
11+
::: faststream.utils.nuid.NUID

faststream/nats/publisher/producer.py

+2-7
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,5 @@
11
import asyncio
2-
from secrets import token_hex
32
from typing import TYPE_CHECKING, Any, Dict, Optional
4-
from uuid import uuid4
53

64
import nats
75
from typing_extensions import override
@@ -71,9 +69,7 @@ async def publish( # type: ignore[override]
7169
if reply_to:
7270
raise WRONG_PUBLISH_ARGS
7371

74-
token = client._nuid.next()
75-
token.extend(token_hex(2).encode())
76-
reply_to = token.decode()
72+
reply_to = client.new_inbox()
7773

7874
future: asyncio.Future[Msg] = asyncio.Future()
7975
sub = await client.subscribe(reply_to, future=future, max_msgs=1)
@@ -148,8 +144,7 @@ async def publish( # type: ignore[override]
148144
if rpc:
149145
if reply_to:
150146
raise WRONG_PUBLISH_ARGS
151-
152-
reply_to = str(uuid4())
147+
reply_to = self._connection._nc.new_inbox()
153148
future: asyncio.Future[Msg] = asyncio.Future()
154149
sub = await self._connection._nc.subscribe(
155150
reply_to, future=future, max_msgs=1

faststream/redis/publisher/producer.py

+4-3
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,4 @@
11
from typing import TYPE_CHECKING, Any, Optional
2-
from uuid import uuid4
32

43
from typing_extensions import override
54

@@ -10,6 +9,7 @@
109
from faststream.redis.parser import RawMessage, RedisPubSubParser
1110
from faststream.redis.schemas import INCORRECT_SETUP_MSG
1211
from faststream.utils.functions import timeout_scope
12+
from faststream.utils.nuid import NUID
1313

1414
if TYPE_CHECKING:
1515
from redis.asyncio.client import PubSub, Redis
@@ -67,8 +67,9 @@ async def publish( # type: ignore[override]
6767
if rpc:
6868
if reply_to:
6969
raise WRONG_PUBLISH_ARGS
70-
71-
reply_to = str(uuid4())
70+
nuid = NUID()
71+
rpc_nuid = str(nuid.next(), "utf-8")
72+
reply_to = rpc_nuid
7273
psub = self._connection.pubsub()
7374
await psub.subscribe(reply_to)
7475

faststream/utils/nuid.py

+68
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
# Copyright 2016-2018 The NATS Authors
2+
# Licensed under the Apache License, Version 2.0 (the "License");
3+
# you may not use this file except in compliance with the License.
4+
# You may obtain a copy of the License at
5+
#
6+
# http://www.apache.org/licenses/LICENSE-2.0
7+
#
8+
# Unless required by applicable law or agreed to in writing, software
9+
# distributed under the License is distributed on an "AS IS" BASIS,
10+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+
# See the License for the specific language governing permissions and
12+
# limitations under the License.
13+
14+
from __future__ import annotations
15+
16+
from random import Random
17+
from secrets import randbelow, token_bytes
18+
from sys import maxsize as max_int
19+
20+
DIGITS = b"0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
21+
BASE = 62
22+
PREFIX_LENGTH = 12
23+
SEQ_LENGTH = 10
24+
MAX_SEQ = 839299365868340224 # BASE**10
25+
MIN_INC = 33
26+
MAX_INC = 333
27+
INC = MAX_INC - MIN_INC
28+
TOTAL_LENGTH = PREFIX_LENGTH + SEQ_LENGTH
29+
30+
31+
class NUID:
32+
"""NUID created is a utility to create a new id.
33+
34+
NUID is an implementation of the approach for fast generation
35+
of unique identifiers used for inboxes in NATS.
36+
"""
37+
38+
def __init__(self) -> None:
39+
self._prand = Random(randbelow(max_int)) # nosec B311
40+
self._seq = self._prand.randint(0, MAX_SEQ)
41+
self._inc = MIN_INC + self._prand.randint(BASE + 1, INC)
42+
self._prefix = bytearray()
43+
self.randomize_prefix()
44+
45+
def next(self) -> bytearray:
46+
"""Next returns the next unique identifier."""
47+
self._seq += self._inc
48+
if self._seq >= MAX_SEQ:
49+
self.randomize_prefix()
50+
self.reset_sequential()
51+
52+
l_seq = self._seq
53+
prefix = self._prefix[:]
54+
suffix = bytearray(SEQ_LENGTH)
55+
for i in reversed(range(SEQ_LENGTH)):
56+
suffix[i] = DIGITS[int(l_seq) % BASE]
57+
l_seq //= BASE
58+
59+
prefix.extend(suffix)
60+
return prefix
61+
62+
def randomize_prefix(self) -> None:
63+
random_bytes = token_bytes(PREFIX_LENGTH)
64+
self._prefix = bytearray(DIGITS[c % BASE] for c in random_bytes)
65+
66+
def reset_sequential(self) -> None:
67+
self._seq = self._prand.randint(0, MAX_SEQ)
68+
self._inc = MIN_INC + self._prand.randint(0, INC)

tests/brokers/nats/test_test_client.py

+12
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,18 @@ def subscriber(m):
8080

8181
assert event.is_set()
8282

83+
@pytest.mark.nats()
84+
async def test_inbox_prefix_with_real(
85+
self,
86+
queue: str,
87+
):
88+
broker = NatsBroker(inbox_prefix="test")
89+
90+
async with TestNatsBroker(broker, with_real=True) as br:
91+
assert br._connection._inbox_prefix == b"test"
92+
assert "test" in str(br._connection.new_inbox())
93+
94+
8395
async def test_respect_middleware(self, queue):
8496
routes = []
8597

0 commit comments

Comments
 (0)