Add configurable timeout for zk_lock retry (#73)

* Add configurable timeout for zk_lock retry * Fix linter * Added exception handling in _zk_flock * Rebased onto lastest main and fixed linting * Restoring _exit_code functioning. Fixing if statement in case of timeout --------- Co-authored-by: Maria Dyuldina <[email protected]> Co-authored-by: Filippo Monari <[email protected]>
yandex · Nov 9, 2023 · 3c28cf5 · 3c28cf5
1 parent e126f91
commit 3c28cf5
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 2 deletions.
diff --git a/ch_backup/config.py b/ch_backup/config.py
@@ -209,6 +209,7 @@ def _as_seconds(t: str) -> int:
         "flock_path": "/tmp/flock.lock",
         "zk_flock_path": "/tmp/zk_flock.lock",
         "exitcode": 0,
+        "lock_timeout": _as_seconds("1 min"),
     },
 }
 

diff --git a/ch_backup/logic/lock_manager.py b/ch_backup/logic/lock_manager.py
@@ -7,6 +7,7 @@
 from types import TracebackType
 from typing import IO, Optional, Type
 
+from kazoo.exceptions import LockTimeout
 from kazoo.recipe.lock import Lock
 
 from ch_backup import logging
@@ -18,7 +19,7 @@ class LockManager:
     Lock manager class
     """
 
-    # pylint: disable=consider-using-with
+    # pylint: disable=consider-using-with,too-many-instance-attributes
 
     _lock_conf: dict
     _process_lockfile_path: str
@@ -41,6 +42,7 @@ def __init__(self, lock_conf: dict, zk_ctl: Optional[ZookeeperCTL]) -> None:
         self._exitcode = lock_conf.get("exitcode")
         self._distributed = zk_ctl is None
         self._disabled = False
+        self._lock_timeout = lock_conf.get("lock_timeout")
 
     def __call__(self, distributed: bool = True, disabled: bool = False):  # type: ignore
         """
@@ -106,7 +108,13 @@ def _zk_flock(self) -> None:
             if not client.exists(self._process_zk_lockfile_path):
                 client.create(self._process_zk_lockfile_path, makepath=True)
             self._zk_lock = client.Lock(self._process_zk_lockfile_path)
-            if not self._zk_lock.acquire(blocking=False):
+            try:
+                _ = self._zk_lock.acquire(blocking=True, timeout=self._lock_timeout)
+                logging.debug("Lock was acquired")
+            except LockTimeout:
+                logging.error(
+                    "Lock was not acquired due to timeout error.", exc_info=True
+                )
                 sys.exit(self._exitcode)
         else:
             raise RuntimeError("ZK flock enabled, but zookeeper is not configured")