Skip to content

Commit

Permalink
feat(recipe): add support for curator SharedCount recipe (#559)
Browse files Browse the repository at this point in the history
* 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
  • Loading branch information
BrianEaton1 authored and StephenSorriaux committed May 7, 2019
1 parent cd49b3f commit 88b657a
Show file tree
Hide file tree
Showing 2 changed files with 48 additions and 5 deletions.
35 changes: 30 additions & 5 deletions kazoo/recipe/counter.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
"""
from kazoo.exceptions import BadVersionError
from kazoo.retry import ForceRetryError

import struct

class Counter(object):
"""Kazoo Counter
Expand All @@ -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.
Expand All @@ -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:
Expand All @@ -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
Expand All @@ -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
Expand Down
18 changes: 18 additions & 0 deletions kazoo/tests/test_counter.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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()
Expand Down

0 comments on commit 88b657a

Please sign in to comment.