From 0160823d084d141be71299c0e51d9dd3c27ede94 Mon Sep 17 00:00:00 2001
From: zsdc <taras@vyos.io>
Date: Wed, 3 Jul 2024 22:45:52 +0300
Subject: locking: T6544: Added locking util

Sometimes we need a reliable way to lock an execution until some other operation
is not done.

This commit introduces locking util, which can be used as a common lock, even
between different processes.

Usage example:

```
from vyos.utils.locking import Lock

lock = Lock('my_lock_id')
lock.acquire(timeout=10)

print(f'Lock acquired: {lock.is_locked}')

lock.release()
```
---
 python/vyos/utils/__init__.py |   1 +
 python/vyos/utils/locking.py  | 115 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 116 insertions(+)
 create mode 100644 python/vyos/utils/locking.py

(limited to 'python/vyos/utils')

diff --git a/python/vyos/utils/__init__.py b/python/vyos/utils/__init__.py
index 90620071b..3759b2125 100644
--- a/python/vyos/utils/__init__.py
+++ b/python/vyos/utils/__init__.py
@@ -25,6 +25,7 @@ from vyos.utils import file
 from vyos.utils import io
 from vyos.utils import kernel
 from vyos.utils import list
+from vyos.utils import locking
 from vyos.utils import misc
 from vyos.utils import network
 from vyos.utils import permission
diff --git a/python/vyos/utils/locking.py b/python/vyos/utils/locking.py
new file mode 100644
index 000000000..63cb1a816
--- /dev/null
+++ b/python/vyos/utils/locking.py
@@ -0,0 +1,115 @@
+# Copyright 2024 VyOS maintainers and contributors <maintainers@vyos.io>
+#
+# This library is free software; you can redistribute it and/or
+# modify it under the terms of the GNU Lesser General Public
+# License as published by the Free Software Foundation; either
+# version 2.1 of the License, or (at your option) any later version.
+#
+# This library is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+# Lesser General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public
+# License along with this library.  If not, see <http://www.gnu.org/licenses/>.
+
+import fcntl
+import re
+import time
+from pathlib import Path
+
+
+class LockTimeoutError(Exception):
+    """Custom exception raised when lock acquisition times out."""
+
+    pass
+
+
+class InvalidLockNameError(Exception):
+    """Custom exception raised when the lock name is invalid."""
+
+    pass
+
+
+class Lock:
+    """Lock class to acquire and release a lock file"""
+
+    def __init__(self, lock_name: str) -> None:
+        """Lock class constructor
+
+        Args:
+            lock_name (str): Name of the lock file
+
+        Raises:
+            InvalidLockNameError: If the lock name is invalid
+        """
+        # Validate lock name
+        if not re.match(r'^[a-zA-Z0-9_\-]+$', lock_name):
+            raise InvalidLockNameError(f'Invalid lock name: {lock_name}')
+
+        self.__lock_dir = Path('/run/vyos/lock')
+        self.__lock_dir.mkdir(parents=True, exist_ok=True)
+
+        self.__lock_file_path: Path = self.__lock_dir / f'{lock_name}.lock'
+        self.__lock_file = None
+
+        self._is_locked = False
+
+    def __del__(self) -> None:
+        """Ensure the lock file is removed when the object is deleted"""
+        self.release()
+
+    @property
+    def is_locked(self) -> bool:
+        """Check if the lock is acquired
+
+        Returns:
+            bool: True if the lock is acquired, False otherwise
+        """
+        return self._is_locked
+
+    def __unlink_lockfile(self) -> None:
+        """Remove the lock file if it is not currently locked."""
+        try:
+            with self.__lock_file_path.open('w') as f:
+                fcntl.flock(f, fcntl.LOCK_EX | fcntl.LOCK_NB)
+                self.__lock_file_path.unlink(missing_ok=True)
+        except IOError:
+            # If we cannot acquire the lock, it means another process has it, so we do nothing.
+            pass
+
+    def acquire(self, timeout: int = 0) -> None:
+        """Acquire a lock file
+
+        Args:
+            timeout (int, optional): A time to wait for lock. Defaults to 0.
+
+        Raises:
+            LockTimeoutError: If lock could not be acquired within timeout
+        """
+        start_time: float = time.time()
+        while True:
+            try:
+                self.__lock_file = self.__lock_file_path.open('w')
+                fcntl.flock(self.__lock_file, fcntl.LOCK_EX | fcntl.LOCK_NB)
+                self._is_locked = True
+                return
+            except IOError:
+                if timeout > 0 and (time.time() - start_time) >= timeout:
+                    if self.__lock_file:
+                        self.__lock_file.close()
+                    raise LockTimeoutError(
+                        f'Could not acquire lock within {timeout} seconds'
+                    )
+                time.sleep(0.1)
+
+    def release(self) -> None:
+        """Release a lock file"""
+        if self.__lock_file and self._is_locked:
+            try:
+                fcntl.flock(self.__lock_file, fcntl.LOCK_UN)
+                self._is_locked = False
+            finally:
+                self.__lock_file.close()
+                self.__lock_file = None
+                self.__unlink_lockfile()
-- 
cgit v1.2.3