diff --git a/pyproject.toml b/pyproject.toml index 7bfa866..edbcb4d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,11 +22,15 @@ classifiers = [ "Programming Language :: Python :: 3.15", "Topic :: Software Development :: Libraries :: Python Modules", ] -dependencies = ["kazoo>=2.10,<3"] +dependencies = [ + "kazoo>=2.10,<3", + "logger>=1.4", +] [dependency-groups] dev = [ "pytest>=8", + "pytest-mock>=3.17.0", "pytest-timeout>=2.3", "ruff==0.17.1", "ty==0.0.90", diff --git a/src/arcus/_cluster/__init__.py b/src/arcus/_cluster/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/arcus/_cluster/discovery.py b/src/arcus/_cluster/discovery.py new file mode 100644 index 0000000..c059683 --- /dev/null +++ b/src/arcus/_cluster/discovery.py @@ -0,0 +1,107 @@ +from collections.abc import Callable + +from kazoo.client import KazooClient +from kazoo.exceptions import KazooException, NoNodeError, OperationTimeoutError +from kazoo.handlers.threading import KazooTimeoutError +from kazoo.protocol.states import WatchedEvent + +from arcus import ArcusConnectionError +from arcus._logging import logger + + +def parse_cache_list(children: list[str]) -> list[str]: + addresses: set[str] = set() + + for child in children: + address = child.partition("-")[0] + host, _, port = address.rpartition(":") + + if ( + not host + or not port + or not port.isascii() + or not port.isdecimal() + or not 1 <= int(port) <= 65535 + ): + logger.warning(f"Ignoring invalid cache node {child}") + continue + + addresses.add(address) + + return sorted(addresses) + + +class ArcusDiscovery: + def __init__( + self, + zk_hosts: str, + service_code: str, + timeout: float = 10.0, + ) -> None: + self._client: KazooClient = KazooClient(hosts=zk_hosts) + self._path: str = f"/arcus/cache_list/{service_code}" + self._timeout: float = timeout + self._closed: bool = False + + # Assigned in start() before registering the watch. + self._on_change: Callable[[list[str]], None] + + def start(self, on_change: Callable[[list[str]], None]) -> list[str]: + self._on_change = on_change + + try: + self._client.start(timeout=self._timeout) + children = self._client.get_children( + path=self._path, + watch=self._watch_children, + ) + return parse_cache_list(children) + except BaseException as error: + self._closed = True + self._cleanup() + + if isinstance(error, (KazooTimeoutError, OperationTimeoutError)): + raise ArcusConnectionError( + "ZooKeeper discovery startup timed out" + ) from error + + if isinstance(error, NoNodeError): + raise ArcusConnectionError( + f"service path not found: {self._path}" + ) from error + + if isinstance(error, (KazooException, OSError)): + raise ArcusConnectionError( + "ZooKeeper discovery startup failed" + ) from error + + raise + + def stop(self) -> None: + self._closed = True + self._cleanup() + + def _watch_children(self, _event: WatchedEvent) -> None: + if self._closed: + return + + try: + children = self._client.get_children( + path=self._path, + watch=self._watch_children, + ) + + self._on_change(parse_cache_list(children)) + except Exception: + logger.exception("Failed to update cache nodes") + + def _cleanup(self) -> None: + try: + self._client.stop() + except Exception: + logger.exception("Failed to stop ZooKeeper") + + try: + self._client.close() + except Exception: + logger.exception("Failed to close ZooKeeper") diff --git a/tests/test_discovery.py b/tests/test_discovery.py new file mode 100644 index 0000000..9f2877b --- /dev/null +++ b/tests/test_discovery.py @@ -0,0 +1,192 @@ +import pytest +from kazoo.exceptions import ConnectionLoss, NoNodeError, OperationTimeoutError +from kazoo.handlers.threading import KazooTimeoutError +from kazoo.protocol.states import EventType, KazooState, WatchedEvent +from pytest_mock import MockerFixture, MockType + +from arcus import ArcusConnectionError +from arcus._cluster.discovery import ArcusDiscovery + +ZK_HOSTS = "127.0.0.1:2181,127.0.0.1:2182,127.0.0.1:2183" +SERVICE_CODE = "test" +SERVICE_PATH = "/arcus/cache_list/test" + +EXPECTED_ADDRESSES = [ + "127.0.0.1:11211", + "127.0.0.1:11212", + "127.0.0.1:11213", +] + + +@pytest.fixture +def zk_client(mocker: MockerFixture) -> MockType: + client_class = mocker.patch( + "arcus._cluster.discovery.KazooClient", + autospec=True, + ) + client = client_class.return_value + client.get_children.return_value = [ + "127.0.0.1:11213-cache3", + "127.0.0.1:11211-cache1", + "127.0.0.1:11212-cache2", + ] + return client + + +@pytest.fixture +def on_change(mocker: MockerFixture) -> MockType: + return mocker.Mock() + + +def test_start_returns_cache_nodes( + zk_client: MockType, + on_change: MockType, +) -> None: + # Given: ZooKeeper provides the initial cache node list. + discovery = ArcusDiscovery(ZK_HOSTS, SERVICE_CODE) + + try: + # When: Discovery starts. + addresses = discovery.start(on_change) + + # Then: It returns sorted addresses and registers a watch. + assert addresses == EXPECTED_ADDRESSES + zk_client.get_children.assert_called_once_with( + path=SERVICE_PATH, + watch=discovery._watch_children, + ) + on_change.assert_not_called() + finally: + discovery.stop() + + +@pytest.mark.parametrize( + ("children", "expected"), + [ + pytest.param( + [ + "127.0.0.1:11214-cache4", + "127.0.0.1:11211-cache1", + ], + ["127.0.0.1:11211", "127.0.0.1:11214"], + id="nodes-changed", + ), + pytest.param([], [], id="all-nodes-removed"), + ], +) +def test_watch_notifies_changes_and_registers_again( + zk_client: MockType, + on_change: MockType, + children: list[str], + expected: list[str], +) -> None: + # Given: Discovery is watching the service and nodes have changed. + discovery = ArcusDiscovery(ZK_HOSTS, SERVICE_CODE) + + try: + discovery.start(on_change) + watch = zk_client.get_children.call_args.kwargs["watch"] + zk_client.get_children.return_value = children + event = WatchedEvent( + EventType.CHILD, + KazooState.CONNECTED, + SERVICE_PATH, + ) + + # When: ZooKeeper delivers a child change event. + watch(event) + + # Then: It notifies the consumer and registers the watch again. + on_change.assert_called_once_with(expected) + assert zk_client.get_children.call_count == 2 + zk_client.get_children.assert_called_with( + path=SERVICE_PATH, + watch=watch, + ) + finally: + discovery.stop() + + +def test_watch_lookup_failure_does_not_notify( + zk_client: MockType, + on_change: MockType, +) -> None: + # Given: Discovery is started, but the next lookup will fail. + discovery = ArcusDiscovery(ZK_HOSTS, SERVICE_CODE) + + try: + discovery.start(on_change) + watch = zk_client.get_children.call_args.kwargs["watch"] + zk_client.get_children.side_effect = ConnectionLoss() + event = WatchedEvent( + EventType.CHILD, + KazooState.CONNECTED, + SERVICE_PATH, + ) + + # When: ZooKeeper delivers a child change event. + watch(event) + + # Then: The failed lookup does not notify the consumer. + assert zk_client.get_children.call_count == 2 + on_change.assert_not_called() + finally: + discovery.stop() + + +@pytest.mark.parametrize( + ("stage", "error"), + [ + pytest.param("connect", ConnectionLoss(), id="connection-failure"), + pytest.param("connect", KazooTimeoutError(), id="connection-timeout"), + pytest.param("lookup", NoNodeError(), id="service-not-found"), + pytest.param("lookup", OperationTimeoutError(), id="lookup-timeout"), + ], +) +def test_start_failure_cleans_up( + zk_client: MockType, + on_change: MockType, + stage: str, + error: Exception, +) -> None: + # Given: Connecting to ZooKeeper or reading the node list will fail. + discovery = ArcusDiscovery(ZK_HOSTS, SERVICE_CODE) + + if stage == "connect": + zk_client.start.side_effect = error + else: + zk_client.get_children.side_effect = error + + # When: Discovery attempts to start. + with pytest.raises(ArcusConnectionError): + discovery.start(on_change) + + # Then: Resources are released without notifying the consumer. + zk_client.stop.assert_called_once_with() + zk_client.close.assert_called_once_with() + on_change.assert_not_called() + + +def test_stop_closes_connection_and_ignores_watch( + zk_client: MockType, + on_change: MockType, +) -> None: + # Given: Discovery is watching the service. + discovery = ArcusDiscovery(ZK_HOSTS, SERVICE_CODE) + discovery.start(on_change) + watch = zk_client.get_children.call_args.kwargs["watch"] + event = WatchedEvent( + EventType.CHILD, + KazooState.CONNECTED, + SERVICE_PATH, + ) + + # When: Discovery is stopped and a late event arrives. + discovery.stop() + watch(event) + + # Then: Resources are released, and the late event is ignored. + zk_client.stop.assert_called_once_with() + zk_client.close.assert_called_once_with() + assert zk_client.get_children.call_count == 1 + on_change.assert_not_called() diff --git a/uv.lock b/uv.lock index b2f79a8..22010ec 100644 --- a/uv.lock +++ b/uv.lock @@ -8,22 +8,28 @@ version = "0.1.0" source = { editable = "." } dependencies = [ { name = "kazoo" }, + { name = "logger" }, ] [package.dev-dependencies] dev = [ { name = "pytest" }, + { name = "pytest-mock" }, { name = "pytest-timeout" }, { name = "ruff" }, { name = "ty" }, ] [package.metadata] -requires-dist = [{ name = "kazoo", specifier = ">=2.10,<3" }] +requires-dist = [ + { name = "kazoo", specifier = ">=2.10,<3" }, + { name = "logger", specifier = ">=1.4" }, +] [package.metadata.requires-dev] dev = [ { name = "pytest", specifier = ">=8" }, + { name = "pytest-mock", specifier = ">=3.17.0" }, { name = "pytest-timeout", specifier = ">=2.3" }, { name = "ruff", specifier = "==0.17.1" }, { name = "ty", specifier = "==0.0.90" }, @@ -56,6 +62,12 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/91/b9/5328ff14f7d637ac90f7144d4aad27b17bdd063325c9fb4b88c4e1f109fe/kazoo-2.11.0-py2.py3-none-any.whl", hash = "sha256:aa6be17159540dae4087661b8d4e24c7b58d4dc26170a2d5b8468b06712d84a9", size = 150943, upload-time = "2026-03-21T19:08:52.575Z" }, ] +[[package]] +name = "logger" +version = "1.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/73/2f/b0d28eaa1e2c1cf64129f8da3fe701888d152677fec708cd0f13e8309e1e/logger-1.4.tar.gz", hash = "sha256:4ecac57133c6376fa215f0fe6b4dc4d60e4d1ad8be005cab4e8a702df682f8b3", size = 1187, upload-time = "2013-10-22T16:49:01.48Z" } + [[package]] name = "packaging" version = "26.3" @@ -99,6 +111,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/24/25/1de2678b631f5a49215c6c96fff41ba892b0a34df68d6d80292b1b48aa7f/pytest-9.1.1-py3-none-any.whl", hash = "sha256:37a86b45efb9a47a61a36449063e8e18d0cab3161329fc099eb21783169c4f0c", size = 386536, upload-time = "2026-06-19T10:58:31.347Z" }, ] +[[package]] +name = "pytest-mock" +version = "3.17.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pytest" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/97/c2/28b9f5b955ff2bcd4727a6f768e2ca9b586cd7c37b95ea5f7b9b873ac7b3/pytest_mock-3.17.0.tar.gz", hash = "sha256:e36576e0985e01027e999891e7a966c67822d40c89ef1b4c8b6e22737e3590b4", size = 43064, upload-time = "2026-10-17T16:21:42.435Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e8/ea/47fbdc8f37410956fd445720e927589b1e2a47772c5f28109f0ce2fa26d9/pytest_mock-3.17.0-py3-none-any.whl", hash = "sha256:1b82aa4b37137248bd687055e109723debeb07b489de1565297dc9589c5f2d6b", size = 11248, upload-time = "2026-10-17T16:21:41.265Z" }, +] + [[package]] name = "pytest-timeout" version = "2.4.0"