From bcb81af661a1de2902092fbb2e74446f06152b3e Mon Sep 17 00:00:00 2001
From: ZipFile <zipfile.d@protonmail.com>
Date: Sat, 17 Aug 2024 14:33:20 +0000
Subject: [PATCH 1/5] Make SonyFlake an iterator

---
 sonyflake/sonyflake.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/sonyflake/sonyflake.py b/sonyflake/sonyflake.py
index c950009..216b37c 100644
--- a/sonyflake/sonyflake.py
+++ b/sonyflake/sonyflake.py
@@ -5,7 +5,7 @@
 from socket import gethostbyname, gethostname
 from threading import Lock
 from time import sleep
-from typing import Any, Callable, Dict, Optional, Union
+from typing import Any, Callable, Dict, Iterator, Optional, Union
 from warnings import warn
 
 BIT_LEN_TIME = 39
@@ -28,13 +28,17 @@ def lower_16bit_private_ip() -> int:
     return (ip_bytes[2] << 8) + ip_bytes[3]
 
 
-class SonyFlake:
+class SonyFlake(Iterator[int]):
     """
     The distributed unique ID generator.
     """
 
+    _now: Callable[[], datetime.datetime]
+    mutex: Lock
     _start_time: int
     _machine_id: int
+    elapsed_time: int
+    sequence: int
 
     __slots__ = (
         "_now",
@@ -148,6 +152,8 @@ def next_id(self) -> int:
                     sleep(self.sleep_time(overtime, self._now()))
             return self.to_id()
 
+    __next__ = next_id
+
     def to_id(self) -> int:
         if self.elapsed_time >= (1 << BIT_LEN_TIME):
             raise TimeoutError("Over the time limit!")

From d265d29f765bc7211c12094f547a6af781a8c508 Mon Sep 17 00:00:00 2001
From: ZipFile <zipfile.d@protonmail.com>
Date: Sat, 17 Aug 2024 14:42:23 +0000
Subject: [PATCH 2/5] Add random_machine_ids()

---
 sonyflake/sonyflake.py  | 19 +++++++++++++++++--
 tests/test_sonyflake.py | 16 +++++++++++++++-
 2 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/sonyflake/sonyflake.py b/sonyflake/sonyflake.py
index 216b37c..c2f4ea1 100644
--- a/sonyflake/sonyflake.py
+++ b/sonyflake/sonyflake.py
@@ -1,11 +1,11 @@
 import datetime
 import ipaddress
 from functools import partial
-from random import randrange
+from random import randrange, sample
 from socket import gethostbyname, gethostname
 from threading import Lock
 from time import sleep
-from typing import Any, Callable, Dict, Iterator, Optional, Union
+from typing import Any, Callable, Dict, Iterator, List, Optional, Union
 from warnings import warn
 
 BIT_LEN_TIME = 39
@@ -19,6 +19,21 @@
 utc_now = partial(datetime.datetime.now, tz=UTC)
 
 
+def random_machine_ids(n: int) -> List[int]:
+    """
+    Returns a list of `n` random machine IDs.
+
+    `n` must be in range (0, 0xFFFF].
+
+    Returned list is sorted in ascending order, without duplicates.
+    """
+
+    if not (0 < n <= MAX_MACHINE_ID):
+        raise ValueError(f"n must be in range (0, {MAX_MACHINE_ID}]")
+
+    return sorted(sample(range(0, MAX_MACHINE_ID + 1), n))
+
+
 def lower_16bit_private_ip() -> int:
     """
     Returns the lower 16 bits of the private IP address.
diff --git a/tests/test_sonyflake.py b/tests/test_sonyflake.py
index b7dcf68..5e1ef65 100644
--- a/tests/test_sonyflake.py
+++ b/tests/test_sonyflake.py
@@ -4,7 +4,7 @@
 from time import sleep
 from unittest import TestCase
 
-from pytest import raises
+from pytest import mark, raises
 
 from sonyflake.sonyflake import (
     BIT_LEN_SEQUENCE,
@@ -12,6 +12,7 @@
     SonyFlake,
     lower_16bit_private_ip,
     random_machine_id,
+    random_machine_ids,
 )
 
 
@@ -110,5 +111,18 @@ def test_random_machine_id() -> None:
     assert random_machine_id()
 
 
+@mark.parametrize("n", [1, 1024, 65535])
+def test_random_machine_ids(n: int) -> None:
+    machine_ids = random_machine_ids(n)
+    assert len(set(machine_ids)) == n
+    assert sorted(machine_ids) == machine_ids
+
+
+@mark.parametrize("n", [0, 65536])
+def test_random_machine_ids_edges(n: int) -> None:
+    with raises(ValueError, match=r"n must be in range \(0, 65535\]"):
+        random_machine_ids(n)
+
+
 def test_lower_16bit_private_ip() -> None:
     assert lower_16bit_private_ip()

From 74f75adbbaca0c4717bed611f697c8f7696245ef Mon Sep 17 00:00:00 2001
From: ZipFile <zipfile.d@protonmail.com>
Date: Sat, 17 Aug 2024 15:22:31 +0000
Subject: [PATCH 3/5] Add RoundRobin wrapper

---
 sonyflake/round_robin.py  | 27 +++++++++++++++++++++++++++
 tests/test_round_robin.py | 18 ++++++++++++++++++
 2 files changed, 45 insertions(+)
 create mode 100644 sonyflake/round_robin.py
 create mode 100644 tests/test_round_robin.py

diff --git a/sonyflake/round_robin.py b/sonyflake/round_robin.py
new file mode 100644
index 0000000..f198fc1
--- /dev/null
+++ b/sonyflake/round_robin.py
@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+from itertools import cycle
+from typing import Iterable, Iterator
+
+
+class RoundRobin(Iterator[int]):
+    """Round-robin iterator for cycling through multiple ID generators.
+
+    Used for generating ids at rate more than 256ids/10msec.
+
+    Example:
+    >>> from sonyflake import RoundRobin, SonyFlake, random_machine_ids
+    >>> sf = RoundRobin([SonyFlake(machine_id=_id) for _id in random_machine_ids(10)])
+    >>> %timeit next(sf)
+    """
+
+    _id_generators: cycle[Iterator[int]]
+    __slots__ = ("_id_generators",)
+
+    def __init__(self, id_generators: Iterable[Iterator[int]]) -> None:
+        self._id_generators = cycle(id_generators)
+
+    def __next__(self) -> int:
+        return next(next(self._id_generators))
+
+    next_id = __next__
diff --git a/tests/test_round_robin.py b/tests/test_round_robin.py
new file mode 100644
index 0000000..25812ec
--- /dev/null
+++ b/tests/test_round_robin.py
@@ -0,0 +1,18 @@
+from sonyflake.round_robin import RoundRobin
+from sonyflake.sonyflake import BIT_LEN_MACHINE_ID, SonyFlake
+
+
+def test_round_robin() -> None:
+    rr = RoundRobin(
+        [
+            SonyFlake(machine_id=0x0000),
+            SonyFlake(machine_id=0x7F7F),
+            SonyFlake(machine_id=0xFFFF),
+        ]
+    )
+
+    assert [next(rr) & ((1 << BIT_LEN_MACHINE_ID) - 1) for _ in range(6)] == [
+        0x0000,
+        0x7F7F,
+        0xFFFF,
+    ] * 2

From 7bf1fecbf2834a5582b990698da85777822b3249 Mon Sep 17 00:00:00 2001
From: ZipFile <zipfile.d@protonmail.com>
Date: Sat, 17 Aug 2024 15:23:26 +0000
Subject: [PATCH 4/5] Expose more in base sonyflake module

---
 sonyflake/__init__.py | 17 +++++++++++++++--
 1 file changed, 15 insertions(+), 2 deletions(-)

diff --git a/sonyflake/__init__.py b/sonyflake/__init__.py
index e1a6c02..9d9a99c 100644
--- a/sonyflake/__init__.py
+++ b/sonyflake/__init__.py
@@ -1,4 +1,17 @@
 from .about import NAME, VERSION, __version__
-from .sonyflake import SonyFlake
+from .round_robin import RoundRobin
+from .sonyflake import (
+    SONYFLAKE_EPOCH,
+    SonyFlake,
+    lower_16bit_private_ip,
+    random_machine_id,
+    random_machine_ids,
+)
 
-__all__ = ["SonyFlake"]
+__all__ = [
+    "RoundRobin",
+    "SonyFlake",
+    "random_machine_id",
+    "random_machine_ids",
+    "lower_16bit_private_ip",
+]

From 66cbfc573df2563b94b140633d72656548d458b7 Mon Sep 17 00:00:00 2001
From: ZipFile <zipfile.d@protonmail.com>
Date: Sat, 17 Aug 2024 15:23:45 +0000
Subject: [PATCH 5/5] Update README.md

---
 README.md | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/README.md b/README.md
index 084ae16..96fe042 100644
--- a/README.md
+++ b/README.md
@@ -38,6 +38,20 @@ custom `machine_id`, `start_time` etc.
 - `machine_id` should be an integer value upto 16-bits, callable or
   `None` (will be used random machine id).
 
+If you need to generate ids at rate more than 256ids/10msec, you can use the `RoundRobin` wrapper over multiple `SonyFlake` instances:
+
+``` python
+from timeit import timeit
+from sonyflake import RoundRobin, SonyFlake, random_machine_ids
+sf = RoundRobin([SonyFlake(machine_id=_id) for _id in random_machine_ids(10)])
+t = timeit(sf.next_id, number=100000)
+print(f"generated 100000 ids in {t:.2f} seconds")
+```
+
+> :warning: This increases the chance of collisions, so be careful when using random machine IDs.
+
+For convenience, both `SonyFlake` and `RoundRobin` implement iterator protocol (`next(sf)`).
+
 ## License
 
 The MIT License (MIT).