Merge branch 'issue325-udp-zk-reuse'

Author: soxofaan

CHANGELOG.md

@@ -28,6 +28,7 @@ without compromising stable operations.
 - Calrissian integration: avoid unnecessary pulls of `alpine` image ([#1132]https://github.com/Open-EO/openeo-geopyspark-driver/issues/1132)
 - Calrissian integration: refactor config to a `CalrissianConfig` sub-config ([#1009](https://github.com/Open-EO/openeo-geopyspark-driver/issues/1009))
 - `save_result`: allow non-string values in `GTiff` `file_metadata` ([#1142](https://github.com/Open-EO/openeo-geopyspark-driver/issues/1142))
+- Add `udp_registry_zookeeper_client_reuse` config for `KazooClient` reuse in `ZooKeeperUserDefinedProcessRepository` ([#1037](https://github.com/Open-EO/openeo-geopyspark-driver/pull/1037))
 
 
 ## 0.64.1

openeogeotrellis/backend.py

@@ -337,11 +337,14 @@ def __init__(
             else ZooKeeperServiceRegistry()
         )
 
-        user_defined_processes = (
-            # TODO #283 #285: eliminate is_ci_context, use some kind of config structure
-            InMemoryUserDefinedProcessRepository() if not use_zookeeper or ConfigParams().is_ci_context
-            else ZooKeeperUserDefinedProcessRepository(hosts=ConfigParams().zookeepernodes)
-        )
+        # TODO #283 #285: eliminate is_ci_context, use some kind of config structure
+        if not use_zookeeper or ConfigParams().is_ci_context:
+            user_defined_processes = InMemoryUserDefinedProcessRepository()
+        else:
+            user_defined_processes = ZooKeeperUserDefinedProcessRepository(
+                hosts=ConfigParams().zookeepernodes,
+                zk_client_reuse=get_backend_config().udp_registry_zookeeper_client_reuse,
+            )
 
         requests_session = requests_with_retry(total=3, backoff_factor=2)
         vault = Vault(ConfigParams().vault_addr, requests_session)

openeogeotrellis/config/config.py

@@ -117,6 +117,8 @@ class GpsBackendConfig(OpenEoBackendConfig):
     use_zk_job_registry: bool = True
     zk_job_registry_max_specification_size: Optional[int] = None
 
+    udp_registry_zookeeper_client_reuse: bool = False
+
     ejr_api: Optional[str] = os.environ.get("OPENEO_EJR_API")
     ejr_backend_id: str = "unknown"
     ejr_credentials_vault_path: Optional[str] = os.environ.get("OPENEO_EJR_CREDENTIALS_VAULT_PATH")

openeogeotrellis/user_defined_process_repository.py

@@ -1,16 +1,15 @@
 import contextlib
 import json
 import logging
-from typing import List, Dict
-from typing import Union
+from typing import Dict, List, Union
 
 from kazoo.client import KazooClient
 from kazoo.exceptions import NodeExistsError, NoNodeError
 from kazoo.handlers.threading import KazooTimeoutError
 from kazoo.retry import KazooRetry
-
-from openeo_driver.backend import UserDefinedProcessMetadata, UserDefinedProcesses
+from openeo_driver.backend import UserDefinedProcesses, UserDefinedProcessMetadata
 from openeo_driver.errors import ProcessGraphNotFoundException
+
 from openeogeotrellis.configparams import ConfigParams
 
 
@@ -20,19 +19,19 @@ class ZooKeeperUserDefinedProcessRepository(UserDefinedProcesses):
 
     _log = logging.getLogger(__name__)
 
-    def __init__(self, hosts: List[str], root: str = "/openeo/udps"):
-        self._hosts = ','.join(hosts)
+    def __init__(self, hosts: List[str], root: str = "/openeo/udps", *, zk_client_reuse: bool = False):
+        self._hosts = ",".join(hosts)
         self._root = root
+        self._zk_client_reuse = zk_client_reuse
+        self._zk_client_cache = None
 
     @staticmethod
     def _serialize(spec: dict) -> bytes:
-        return json.dumps({
-            'specification': spec
-        }).encode()
+        return json.dumps({"specification": spec}).encode("utf8")
 
     @staticmethod
     def _deserialize(data: bytes) -> dict:
-        return json.loads(data.decode())
+        return json.loads(data.decode("utf8"))
 
     def save(self, user_id: str, process_id: str, spec: dict) -> None:
         with self._zk_client() as zk:
@@ -82,16 +81,33 @@ def delete(self, user_id: str, process_id: str) -> None:
 
     @contextlib.contextmanager
     def _zk_client(self):
-        kz_retry = KazooRetry(max_tries=10, delay=0.5, backoff=2)
-        zk = KazooClient(hosts=self._hosts,connection_retry=kz_retry,
-                 command_retry=kz_retry, timeout=3.0)
-        zk.start(timeout=15.0)
+        create_new_client = (not self._zk_client_reuse) or (self._zk_client_cache is None)
+        if create_new_client:
+            kz_retry = KazooRetry(max_tries=10, delay=0.5, backoff=2)
+            zk = KazooClient(
+                hosts=self._hosts,
+                connection_retry=kz_retry,
+                command_retry=kz_retry,
+                timeout=3.0,
+            )
+            zk.start(timeout=15.0)
+            if self._zk_client_reuse:
+                self._zk_client_cache = zk
+        else:
+            zk = self._zk_client_cache
 
         try:
             yield zk
         finally:
-            zk.stop()
-            zk.close()
+            if not self._zk_client_reuse:
+                zk.stop()
+                zk.close()
+
+    def __del__(self):
+        if self._zk_client_cache:
+            self._zk_client_cache.stop()
+            self._zk_client_cache.close()
+            self._zk_client_cache = None
 
 
 class InMemoryUserDefinedProcessRepository(UserDefinedProcesses):
@@ -122,7 +138,7 @@ def delete(self, user_id: str, process_id: str) -> None:
 
 
 def main():
-    repo = ZooKeeperUserDefinedProcessRepository(hosts=ConfigParams().zookeepernodes)
+    repo = ZooKeeperUserDefinedProcessRepository(hosts=ConfigParams().zookeepernodes, zk_client_reuse=True)
 
     user_id = 'vdboschj'
     process_graph_id = 'evi'

tests/test_user_defined_process_repository.py

@@ -1,5 +1,6 @@
 import collections
 import json
+from typing import Iterator
 from unittest import mock
 
 import pytest
@@ -16,7 +17,6 @@
     "add": {"process_id": "add", "arguments": {"x": 2, "y": 3}, "result": True}
 }
 
-PG_2PLUS3_JSON = json.dumps(PG_2PLUS3)
 
 
 class TestInMemoryUserDefinedProcessRepository:
@@ -73,10 +73,28 @@ def udp_db(self) -> ZooKeeperUserDefinedProcessRepository:
         return udp_db
 
     @pytest.fixture
-    def zk(self) -> mock.MagicMock:
+    def kazoo_client_factory(self) -> Iterator[mock.MagicMock]:
+        """
+        Inspectable mock for `KazooClient` used by `ZooKeeperUserDefinedProcessRepository` to
+        inspect the calls made to it (e.g. check client reuse).
+        """
+        with mock.patch("openeogeotrellis.user_defined_process_repository.KazooClient") as KazooClient:
+            yield KazooClient
+
+    @pytest.fixture
+    def zk(self, kazoo_client_factory) -> mock.MagicMock:
         """Fixture for `KazooClient instance used by ZooKeeperUserDefinedProcessRepository """
-        with mock.patch('openeogeotrellis.user_defined_process_repository.KazooClient') as KazooClient:
-            yield KazooClient()
+        return kazoo_client_factory()
+
+    def _build_get_return_value(self, *, process_id="evi", pg=PG_2PLUS3):
+        """Helper to build a return value for a mocked `zk.get`"""
+        data = ZooKeeperUserDefinedProcessRepository._serialize(
+            spec={
+                "id": process_id,
+                "process_graph": pg,
+            }
+        )
+        return (data, self.ZnodeStat(version=123))
 
     def test_save_create(self, udp_db, zk):
         udp_db.save(user_id="john", process_id="evi", spec={"id": "evi", "process_graph": PG_2PLUS3})
@@ -95,8 +113,7 @@ def test_get_miss(self, udp_db, zk):
         assert res is None
 
     def test_get(self, udp_db, zk):
-        data = ('{"specification": {"id": "evi", "process_graph": %s}}' % PG_2PLUS3_JSON).encode("utf-8")
-        zk.get.return_value = (data, self.ZnodeStat(version=123))
+        zk.get.return_value = self._build_get_return_value(process_id="evi", pg=PG_2PLUS3)
         res = udp_db.get(user_id="john", process_id="evi")
         assert res == UserDefinedProcessMetadata(process_graph=PG_2PLUS3, id="evi")
 
@@ -114,3 +131,41 @@ def test_delete_miss(self, udp_db, zk):
         zk.delete.side_effect = NoNodeError()
         with pytest.raises(ProcessGraphNotFoundException):
             udp_db.delete(user_id="john", process_id="evi1")
+
+    @pytest.mark.parametrize(
+        ["zk_client_reuse", "expected_create_call_counts", "expected_shutdown_call_counts"],
+        [
+            (False, [0, 1, 2, 3], [3, 3]),
+            (True, [0, 1, 1, 1], [0, 1]),
+        ],
+    )
+    def test_client_reuse(
+        self, kazoo_client_factory, zk_client_reuse, expected_create_call_counts, expected_shutdown_call_counts
+    ):
+        udp_db = ZooKeeperUserDefinedProcessRepository(hosts=["zk1.test", "zk2.test"], zk_client_reuse=zk_client_reuse)
+
+        # Check reuse
+        create_call_counts = []
+        create_call_counts.append(kazoo_client_factory.call_count)
+        udp_db.save(user_id="john", process_id="aaa", spec={"id": "aaa", "process_graph": PG_2PLUS3})
+        create_call_counts.append(kazoo_client_factory.call_count)
+        udp_db.save(user_id="john", process_id="bbb", spec={"id": "bbb", "process_graph": PG_2PLUS3})
+        create_call_counts.append(kazoo_client_factory.call_count)
+
+        kazoo_client_factory.return_value.get.return_value = self._build_get_return_value(process_id="ccc")
+        _ = udp_db.get(user_id="john", process_id="ccc")
+        create_call_counts.append(kazoo_client_factory.call_count)
+
+        assert create_call_counts == expected_create_call_counts
+
+        # Check stop/close operations on `del`
+        stop_call_counts = []
+        close_call_counts = []
+        stop_call_counts.append(kazoo_client_factory.return_value.stop.call_count)
+        close_call_counts.append(kazoo_client_factory.return_value.close.call_count)
+        del udp_db
+        stop_call_counts.append(kazoo_client_factory.return_value.stop.call_count)
+        close_call_counts.append(kazoo_client_factory.return_value.close.call_count)
+
+        assert stop_call_counts == expected_shutdown_call_counts
+        assert close_call_counts == expected_shutdown_call_counts