From 88b657a0977161f3815657878ba48f82a97a3846 Mon Sep 17 00:00:00 2001 From: BrianEaton1 Date: Tue, 7 May 2019 09:38:58 -0400 Subject: [PATCH] feat(recipe): add support for curator SharedCount recipe (#559) * feat(recipe): add support for curator SharedCount recipe This feature allows Java clients using curator's SharedCount recipe and python clients using kazoo's Counter recipe to read and write from the same path without receiving type errors. example use: counter = zk.Counter("/curator", support_curator=True) counter += 2 counter -= 1 counter.value == 1 counter.pre_value == 2 counter.post_value == 1 Closes #558 --- kazoo/recipe/counter.py | 35 ++++++++++++++++++++++++++++++----- kazoo/tests/test_counter.py | 18 ++++++++++++++++++ 2 files changed, 48 insertions(+), 5 deletions(-) diff --git a/kazoo/recipe/counter.py b/kazoo/recipe/counter.py index b728bc2a..9e688491 100644 --- a/kazoo/recipe/counter.py +++ b/kazoo/recipe/counter.py @@ -6,7 +6,7 @@ """ from kazoo.exceptions import BadVersionError from kazoo.retry import ForceRetryError - +import struct class Counter(object): """Kazoo Counter @@ -19,6 +19,12 @@ class Counter(object): `type(counter.default)(value)` both using an ascii encoding. As such other data types might be used for the counter value. + If you would like to support clients updating the same znode path using + either kazoo's counter recipe or curator's SharedCount recipe, you will + need to enable the support_curator flag. This flag limits + support to integers only and does not use ascii encoding as described + above. + Counter changes can raise :class:`~kazoo.exceptions.BadVersionError` if the retry policy wasn't able to apply a change. @@ -42,22 +48,35 @@ class Counter(object): counter.pre_value == 1.0 counter.post_value == 3.0 + counter = zk.Counter("/curator", support_curator=True) + counter += 2 + counter -= 1 + counter.value == 1 + counter.pre_value == 2 + counter.post_value == 1 + """ - def __init__(self, client, path, default=0): + def __init__(self, client, path, default=0, support_curator=False): """Create a Kazoo Counter :param client: A :class:`~kazoo.client.KazooClient` instance. :param path: The counter path to use. - :param default: The default value. + :param default: The default value to use for new counter paths. + :param support_curator: Enable if support for curator's SharedCount + recipe is desired. """ self.client = client self.path = path self.default = default self.default_type = type(default) + self.support_curator = support_curator self._ensured_path = False self.pre_value = None self.post_value = None + if self.support_curator and not isinstance(self.default, int): + raise TypeError("when support_curator is enabled the default " + "type must be an int") def _ensure_node(self): if not self._ensured_path: @@ -68,7 +87,10 @@ def _ensure_node(self): def _value(self): self._ensure_node() old, stat = self.client.get(self.path) - old = old.decode('ascii') if old != b'' else self.default + if self.support_curator: + old = struct.unpack(">i", old)[0] if old != b'' else self.default + else: + old = old.decode('ascii') if old != b'' else self.default version = stat.version data = self.default_type(old) return data, version @@ -86,7 +108,10 @@ def _change(self, value): def _inner_change(self, value): self.pre_value, version = self._value() post_value = self.pre_value + value - data = repr(post_value).encode('ascii') + if self.support_curator: + data = struct.pack(">i", post_value) + else: + data = repr(post_value).encode('ascii') try: self.client.set(self.path, data, version=version) except BadVersionError: # pragma: nocover diff --git a/kazoo/tests/test_counter.py b/kazoo/tests/test_counter.py index 50095a25..47e077e4 100644 --- a/kazoo/tests/test_counter.py +++ b/kazoo/tests/test_counter.py @@ -21,6 +21,22 @@ def test_int_counter(self): counter - 1 eq_(counter.value, -1) + def test_int_curator_counter(self): + counter = self._makeOne(support_curator=True) + eq_(counter.value, 0) + counter += 2 + counter + 1 + eq_(counter.value, 3) + counter -= 3 + counter - 1 + eq_(counter.value, -1) + counter += 1 + counter += 2147483647 + eq_(counter.value, 2147483647) + counter -= 2147483647 + counter -= 2147483647 + eq_(counter.value, -2147483647) + def test_float_counter(self): counter = self._makeOne(default=0.0) eq_(counter.value, 0.0) @@ -33,6 +49,8 @@ def test_errors(self): counter = self._makeOne() self.assertRaises(TypeError, counter.__add__, 2.1) self.assertRaises(TypeError, counter.__add__, b"a") + with self.assertRaises(TypeError): + counter = self._makeOne(default=0.0, support_curator=True) def test_pre_post_values(self): counter = self._makeOne()