lockable_resources: add LockedResourceReservation context manager and various ResourceSelectors

Author: cdleonard

jenkinsapi/lockable_resources.py

@@ -1,10 +1,20 @@
 import logging
-from typing import TYPE_CHECKING, Dict, List, Mapping, Optional, TypedDict
+import time
+from typing import (
+    TYPE_CHECKING,
+    Dict,
+    Iterator,
+    List,
+    Mapping,
+    Optional,
+    TypedDict,
+)
+
+from requests import Response
 
 from jenkinsapi.custom_exceptions import JenkinsAPIException
 from jenkinsapi.jenkinsbase import JenkinsBase
 from jenkinsapi.utils.requester import Requester
-from requests import Response
 
 if TYPE_CHECKING:
     from jenkinsapi.jenkins import Jenkins
@@ -184,3 +194,205 @@ def reserve(self, name: str) -> None:
 
     def unreserve(self, name: str) -> None:
         self._make_resource_request("unreserve", name)
+
+    def try_reserve(
+        self,
+        selector: "ResourceSelector",
+    ) -> Optional[str]:
+        """
+        Try to reserve a resource that matches the given condition
+
+        :return: the name of the reserved resource on success
+        :return: None if all resources are busy
+        """
+        for resource_name in selector.select(self):
+            resource = self[resource_name]
+            # if server reported that the resource is not free
+            # don't try to reserve it
+            if not resource.is_free():
+                continue
+            # if server reported that the resource is free
+            # it might have been reserved since the last poll
+            try:
+                resource.reserve()
+            except ResourceLockedError:
+                continue
+            return resource.name
+        return None
+
+    DEFAULT_WAIT_SLEEP_PERIOD = 5
+    DEFAULT_WAIT_TIMEOUT_PERIOD = 3600
+
+    def wait_reserve(
+        self,
+        selector: "ResourceSelector",
+        sleep_period: float = DEFAULT_WAIT_SLEEP_PERIOD,
+        timeout: float = DEFAULT_WAIT_TIMEOUT_PERIOD,
+    ) -> str:
+        """
+        Reserve a resource that matches the given condition
+
+        :return: the name of the reserved resource
+        :raise TimeoutError: if no resource could be reserved
+        """
+        start_time = time.time()
+        while True:
+            result = self.try_reserve(selector)
+            if result is not None:
+                return result
+            if time.time() - start_time > timeout:
+                raise TimeoutError("Timed out waiting for resource")
+            logger.info(
+                "No free resources matching %r, sleep %.3f seconds",
+                selector,
+                sleep_period,
+            )
+            self.poll()
+            time.sleep(sleep_period)
+
+    def reservation_by_label(
+        self,
+        label: str,
+        sleep_period: float = DEFAULT_WAIT_SLEEP_PERIOD,
+        timeout: float = DEFAULT_WAIT_TIMEOUT_PERIOD,
+    ) -> "LockedResourceReservation":
+        return LockedResourceReservation(
+            self,
+            ResourceLabelSelector(label),
+            sleep_period=sleep_period,
+            timeout=timeout,
+        )
+
+    def reservation_by_name(
+        self,
+        name: str,
+        sleep_period: float = DEFAULT_WAIT_SLEEP_PERIOD,
+        timeout: float = DEFAULT_WAIT_TIMEOUT_PERIOD,
+    ) -> "LockedResourceReservation":
+        return LockedResourceReservation(
+            self,
+            ResourceNameSelector(name),
+            sleep_period=sleep_period,
+            timeout=timeout,
+        )
+
+    def reservation_by_name_list(
+        self,
+        name_list: List[str],
+        sleep_period: float = DEFAULT_WAIT_SLEEP_PERIOD,
+        timeout: float = DEFAULT_WAIT_TIMEOUT_PERIOD,
+    ) -> "LockedResourceReservation":
+        return LockedResourceReservation(
+            self,
+            ResourceNameListSelector(name_list),
+            sleep_period=sleep_period,
+            timeout=timeout,
+        )
+
+
+class ResourceSelector:
+    """Base class for which iterates acceptable resources for a reservation"""
+
+    def select(self, lockable_resources: LockableResources) -> Iterator[str]:
+        raise NotImplementedError("Subclasses must implement __call__ method")
+
+
+class ResourceNameSelector(ResourceSelector):
+    def __init__(self, name: str):
+        self.name = name
+
+    def select(self, lockable_resources: LockableResources) -> Iterator[str]:
+        yield self.name
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({self.name!r})"
+
+
+class ResourceNameListSelector(ResourceSelector):
+    def __init__(self, name_list: List[str]):
+        self.name_list = name_list
+
+    def select(self, lockable_resources: LockableResources) -> Iterator[str]:
+        return iter(self.name_list)
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({self.name_list!r})"
+
+
+class ResourceLabelSelector(ResourceSelector):
+    def __init__(self, label: str):
+        self.label = label
+
+    def __call__(self, resource: LockableResource) -> bool:
+        return self.label in resource.data.get("labels", [])
+
+    def select(self, lockable_resources: LockableResources) -> Iterator[str]:
+        for resource in lockable_resources.values():
+            if self.label in resource.data["labelsAsList"]:
+                yield resource.name
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({self.label!r})"
+
+
+class LockedResourceReservation:
+    """
+    Context manager for locking a Jenkins resource
+
+    Creating this object does not lock the resource, it is only locked and
+    unlocked on :meth:`__enter__` and :meth:`__exit__` methods.
+
+    Example::
+
+        reservation: LockedResourceReservation = init_reservation()
+        # .. possibly much later ...
+        print("Resource will be locked ...")
+        with reservation as locked_resource:
+            name = locked_resource.locked_resource_name
+            print(f"Resource currently locked: {name}")
+        print("Resource no longer locked")
+
+    If resources are busy this will retry until it will eventually succeed or
+    raise a `TimeoutError`.
+    """
+
+    _locked_resource_name: Optional[str] = None
+    locked: bool = False
+
+    def __init__(
+        self,
+        api: LockableResources,
+        selector: ResourceSelector,
+        sleep_period: float = LockableResources.DEFAULT_WAIT_SLEEP_PERIOD,
+        timeout: float = LockableResources.DEFAULT_WAIT_TIMEOUT_PERIOD,
+    ):
+        self.api = api
+        self.selector = selector
+        self.sleep_period = sleep_period
+        self.timeout = timeout
+
+    @property
+    def locked_resource_name(self) -> str:
+        """
+        Return the name of the locked resource
+
+        This throws an error if the resource is not currently locked.
+        """
+        if self._locked_resource_name is None:
+            raise RuntimeError("Resource not locked")
+        return self._locked_resource_name
+
+    def __enter__(self) -> "LockedResourceReservation":
+        """Acquire a lock for the specified label."""
+        if self._locked_resource_name is not None:
+            raise RuntimeError("Lock already acquired")
+        self._locked_resource_name = self.api.wait_reserve(
+            self.selector, sleep_period=self.sleep_period, timeout=self.timeout
+        )
+        self.locked = True
+        return self
+
+    def __exit__(self, *a) -> None:
+        if self._locked_resource_name is not None:
+            self.api.unreserve(self._locked_resource_name)
+            self._locked_resource_name = None

jenkinsapi_tests/systests/test_lockable_resources.py

@@ -37,6 +37,16 @@ def test_lock_name() -> str:
     return "locktest"
 
 
+@pytest.fixture
+def test_lock_name2() -> str:
+    return "locktest2"
+
+
+@pytest.fixture
+def test_lock_label() -> str:
+    return "locktest"
+
+
 @pytest.fixture(scope="function")
 def lockable_resources(
     jenkins_admin_admin: Jenkins,
@@ -101,3 +111,50 @@ def test_reserve_unreserve_nopoll(
     assert lockable_resources.is_reserved(rn) is True
     lockable_resources.poll()
     assert lockable_resources.is_reserved(rn) is False
+
+
+def test_reservation_by_name(
+    lockable_resources: LockableResources,
+    test_lock_name: str,
+):
+    reservation = lockable_resources.reservation_by_name(test_lock_name)
+    assert lockable_resources.is_free(test_lock_name)
+    with reservation:
+        assert reservation.locked_resource_name == test_lock_name
+        assert lockable_resources.is_free(test_lock_name) is False
+    assert lockable_resources.is_free(test_lock_name)
+    name = None
+    with pytest.raises(RuntimeError):
+        name = reservation.locked_resource_name
+    assert name is None
+
+
+def test_reservation_by_name_list(
+    lockable_resources: LockableResources,
+    test_lock_name: str,
+    test_lock_name2: str,
+):
+    name_list = [test_lock_name, test_lock_name2]
+    r1 = lockable_resources.reservation_by_name_list(name_list)
+    assert lockable_resources.is_free(name_list[0])
+    assert lockable_resources.is_free(name_list[1])
+    with lockable_resources.reservation_by_name_list(name_list) as r1:
+        assert r1.locked_resource_name == name_list[0]
+        with lockable_resources.reservation_by_name_list(name_list) as r2:
+            assert r2.locked_resource_name == name_list[1]
+            assert lockable_resources.is_free(name_list[1]) is False
+        assert lockable_resources.is_free(name_list[1])
+    assert lockable_resources.is_free(name_list[0])
+    assert lockable_resources.is_free(name_list[1])
+
+
+def test_reservation_by_label(
+    lockable_resources: LockableResources,
+    test_lock_label: str,
+):
+    res = lockable_resources.reservation_by_label(test_lock_label)
+    with res:
+        locked_resource = lockable_resources[res.locked_resource_name]
+        assert locked_resource.is_free() is False
+        assert test_lock_label in locked_resource.data["labelsAsList"]
+    assert locked_resource.is_free() is True