add round-robin utilities

.gitignore

@@ -0,0 +1,6 @@
+*.pyc
+.DS_Store
+.pyre
+.pyre_configuration
+.watchmanconfig
+__pycache__

README.md

@@ -0,0 +1,17 @@
+This is rather small collection of round robin utilites
+
+```python
+>>> import roundrobin
+>>> get_roundrobin = roundrobin.basic(["A", "B", "C"])
+>>> ''.join([get_roundrobin() for _ in range(7)])
+'ABCABCA'
+>>> # weighted round-robin balancing algorithm as seen in LVS
+>>> get_weighted = roundrobin.weighted([("A", 5), ("B", 1), ("C", 1)])
+>>> ''.join([get_weighted() for _ in range(7)])
+'AAAAABC'
+>>> # smooth weighted round-robin balancing algorithm as seen in Nginx
+>>> get_weighted_smooth = roundrobin.smooth([("A", 5), ("B", 1), ("C", 1)])
+>>> ''.join([get_weighted_smooth() for _ in range(7)])
+'AABACAA'
+```
+

roundrobin/__init__.py

@@ -0,0 +1,5 @@
+from roundrobin.basic_rr import basic
+from roundrobin.smooth_rr import smooth
+from roundrobin.weighted_rr import weighted
+
+__ALL__ = [basic, weighted, smooth]

roundrobin/basic_rr.py

@@ -0,0 +1,11 @@
+from itertools import cycle, islice
+from typing import Any, Callable
+
+
+def basic(dataset: [Any]) -> Callable[[], Any]:
+    iterator = cycle(dataset)
+
+    def get_next():
+        return next(iterator)
+
+    return get_next

roundrobin/smooth_rr.py

@@ -0,0 +1,41 @@
+from typing import Callable, Optional, Tuple
+
+
+def smooth(dataset: [Tuple[str, int]]) -> Callable[[], Optional[str]]:
+    dataset_length = len(dataset)
+    dataset_extra_weights = [ItemWeight(*x) for x in dataset]
+
+    def get_next() -> Optional[str]:
+        if dataset_length == 0:
+            return None
+        if dataset_length == 1:
+            return dataset[0][0]
+
+        total_weight = 0
+        result = None
+        for extra in dataset_extra_weights:
+            extra.current_weight += extra.effective_weight
+            total_weight += extra.effective_weight
+            if extra.effective_weight < extra.weight:
+                extra.effective_weight += 1
+            if not result or result.current_weight < extra.current_weight:
+                result = extra
+        if result:
+            result.current_weight -= total_weight
+            return result.key
+        return None
+
+    return get_next
+
+
+class ItemWeight:
+    key: str
+    weight: int
+    current_weight: int
+    effective_weight: int
+
+    def __init__(self, key, weight):
+        self.key = key
+        self.weight = weight
+        self.current_weight = 0
+        self.effective_weight = weight

roundrobin/weighted_rr.py

@@ -0,0 +1,31 @@
+import math
+from typing import Callable, Optional, Tuple
+
+
+def weighted(dataset: [Tuple[str, int]]) -> Callable[[], Optional[str]]:
+    current_index = -1
+    current_weight = 0
+    dataset_length = len(dataset)
+    dataset_max_weight = 0
+    dataset_gcd_weight = 0
+
+    for _, weight in dataset:
+        if dataset_max_weight < weight:
+            dataset_max_weight = weight
+        dataset_gcd_weight = math.gcd(dataset_gcd_weight, weight)
+
+    def get_next() -> Optional[str]:
+        nonlocal current_index
+        nonlocal current_weight
+        while True:
+            current_index = (current_index + 1) % dataset_length
+            if current_index == 0:
+                current_weight = current_weight - dataset_gcd_weight
+                if current_weight <= 0:
+                    current_weight = dataset_max_weight
+                    if current_weight == 0:
+                        return None
+            if dataset[current_index][1] >= current_weight:
+                return dataset[current_index][0]
+
+    return get_next

setup.py

@@ -0,0 +1,24 @@
+import pathlib
+
+import setuptools
+
+HERE = pathlib.Path(__file__).parent
+README = (HERE / "README.md").read_text()
+
+setuptools.setup(
+    name="roundrobin",
+    version="0.0.1",
+    description="Collection of roundrobin utilities",
+    long_description=README,
+    long_description_content_type="text/markdown",
+    url="https://github.com/linnik/roundrobin",
+    author="Vyacheslav Linnik",
+    author_email="[email protected]",
+    license="MIT",
+    classifiers=[
+        "License :: OSI Approved :: MIT License",
+        "Programming Language :: Python :: 3",
+        "Programming Language :: Python :: 3.5",
+    ],
+    packages=setuptools.find_packages(),
+)

test.py

@@ -0,0 +1,48 @@
+import unittest
+
+import roundrobin
+
+
+class WeightCase(unittest.TestCase):
+
+    def test_basic(self):
+        data = ["A", "B", "C"]
+        get_next = roundrobin.basic(data)
+        result = ''.join([get_next() for _ in range(7)])
+        self.assertEqual(result, 'ABCABCA')
+
+    def test_smooth(self):
+        data = [("A", 5), ("B", 1), ("C", 1)]
+        get_next = roundrobin.smooth(data)
+        result = ''.join([get_next() for _ in range(7)])
+        self.assertEqual(result, 'AABACAA')
+
+    def test_weighted_not_smooth(self):
+        data = [("A", 5), ("B", 1), ("C", 1)]
+        get_next = roundrobin.weighted(data)
+        result = ''.join([get_next() for _ in range(7)])
+        self.assertEqual(result, 'AAAAABC')
+
+    def test_weighted(self):
+        data = [("A", 50), ("B", 35), ("C", 15)]
+        get_next = roundrobin.weighted(data)
+
+        cnt = {}
+        for _ in range(100):
+            key = get_next()
+            cnt[key] = cnt.get(key, 0) + 1
+        self.assertTrue(cnt['A'] == 50)
+        self.assertTrue(cnt['B'] == 35)
+        self.assertTrue(cnt['C'] == 15)
+
+        cnt = {}
+        for _ in range(200):
+            key = get_next()
+            cnt[key] = cnt.get(key, 0) + 1
+        self.assertTrue(cnt['A'] == 100)
+        self.assertTrue(cnt['B'] == 70)
+        self.assertTrue(cnt['C'] == 30)
+
+
+if __name__ == "__main__":
+    unittest.main()