lockable_resources: add LockedResourceReservation context manager and various ResourceSelectors

cdleonard · cdleonard · commit 817fbabb8c83 · 2025-08-25T19:24:49.000+03:00
diff --git a/jenkinsapi/lockable_resources.py b/jenkinsapi/lockable_resources.py
@@ -1,9 +1,20 @@
+from abc import ABC, abstractmethod
 import logging
-from typing import TYPE_CHECKING, Dict, List, Mapping, Optional, TypedDict
+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 requests import Response
+from jenkinsapi.utils.retry import RetryConfig, SimpleRetryConfig
 
 if TYPE_CHECKING:
     from jenkinsapi.jenkins import Jenkins
@@ -86,6 +97,20 @@ class ResourceLockedError(JenkinsAPIException):
     pass
 
 
+class ResourceReservationTimeoutError(JenkinsAPIException, TimeoutError):
+    """Raised when resource reservation times out"""
+
+    pass
+
+
+DEFAULT_WAIT_SLEEP_PERIOD = 5
+DEFAULT_WAIT_TIMEOUT_PERIOD = 3600
+DEFAULT_RETRY_CONFIG = SimpleRetryConfig(
+    sleep_period=DEFAULT_WAIT_SLEEP_PERIOD,
+    timeout=DEFAULT_WAIT_TIMEOUT_PERIOD,
+)
+
+
 class LockableResources(JenkinsBase, Mapping[str, LockableResource]):
     """Object representation of the lockable resource jenkins API"""
 
@@ -179,3 +204,201 @@ 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
+
+    def wait_reserve(
+        self,
+        selector: "ResourceSelector",
+        retry: RetryConfig = DEFAULT_RETRY_CONFIG,
+    ) -> str:
+        """
+        Wait for a resource that matches the given condition to become available
+
+        :return: the name of the reserved resource on success
+        :raise ResourceReservationTimeoutError: if no matching resources are found during the timeout period.
+        """
+        retry_state = retry.begin()
+        while True:
+            result = self.try_reserve(selector)
+            if result is not None:
+                return result
+            try:
+                retry_state.check_retry()
+            except TimeoutError as err:
+                raise ResourceReservationTimeoutError(
+                    f"Timeout waiting for a resource matching {selector} after {retry}"
+                ) from err
+            logger.info("No free resources matching %r, retry", selector)
+            self.poll()
+
+    def reservation_by_label(
+        self,
+        label: str,
+        retry: RetryConfig = DEFAULT_RETRY_CONFIG,
+    ) -> "LockedResourceReservation":
+        return LockedResourceReservation(
+            self,
+            ResourceLabelSelector(label),
+            retry=retry,
+        )
+
+    def reservation_by_name(
+        self,
+        name: str,
+        retry: RetryConfig = DEFAULT_RETRY_CONFIG,
+    ) -> "LockedResourceReservation":
+        return LockedResourceReservation(
+            self,
+            ResourceNameSelector(name),
+            retry=retry,
+        )
+
+    def reservation_by_name_list(
+        self,
+        name_list: List[str],
+        retry: RetryConfig = DEFAULT_RETRY_CONFIG,
+    ) -> "LockedResourceReservation":
+        return LockedResourceReservation(
+            self,
+            ResourceNameListSelector(name_list),
+            retry=retry,
+        )
+
+
+class ResourceSelector(ABC):
+    """Base class for which iterates acceptable resources for a reservation"""
+
+    @abstractmethod
+    def select(self, lockable_resources: LockableResources) -> Iterator[str]:
+        """Iterate acceptable resource names"""
+        pass
+
+
+class ResourceNameSelector(ResourceSelector):
+    """Implementation of :py:class:`ResourceSelector` that selects a single resource by name"""
+
+    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):
+    """Implementation of :py:class:`ResourceSelector` that selects from a list of resources"""
+
+    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):
+    """Implementation of :py:class:`ResourceSelector` that selects any resources with a given jenkins label"""
+
+    def __init__(self, label: str):
+        self.label = label
+
+    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 time out.
+
+    :raise ResourceReservationTimeoutError: if reservation process times out
+    """
+
+    _locked_resource_name: Optional[str] = None
+    retry: RetryConfig
+
+    def __init__(
+        self,
+        api: LockableResources,
+        selector: ResourceSelector,
+        retry: RetryConfig = DEFAULT_RETRY_CONFIG,
+    ):
+        self.api = api
+        self.selector = selector
+        self.retry = retry
+
+    def is_active(self) -> bool:
+        """Check if the resource is currently locked"""
+        return self._locked_resource_name is not None
+
+    @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, retry=self.retry
+        )
+        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
diff --git a/jenkinsapi_tests/systests/test_lockable_resources.py b/jenkinsapi_tests/systests/test_lockable_resources.py
@@ -1,12 +1,16 @@
+from contextlib import ExitStack
+from unittest import mock
 import pytest
 
 from jenkinsapi.jenkins import Jenkins
 from jenkinsapi.lockable_resources import (
     LockableResource,
     LockableResources,
     ResourceLockedError,
+    ResourceReservationTimeoutError,
 )
 from jenkinsapi.utils.jenkins_launcher import JenkinsLancher
+from jenkinsapi.utils.retry import SimpleRetryConfig
 
 GROOVY_SCRIPT_INIT_TEST_RESOURCES = """
 import org.jenkins.plugins.lockableresources.*
@@ -37,6 +41,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 +115,95 @@ 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 reservation.is_active()
+    assert lockable_resources.is_free(test_lock_name)
+    assert reservation.is_active() is False
+    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
+
+
+def test_custom_retry(
+    lockable_resources: LockableResources,
+    test_lock_name: str,
+):
+    with ExitStack() as exit_stack:
+        exit_stack.enter_context(
+            mock.patch(
+                "time.monotonic",
+                side_effect=range(1000, 10000),
+            )
+        )
+        mock_time_sleep = exit_stack.enter_context(
+            mock.patch("time.sleep"),
+        )
+        mock_try_reserve = exit_stack.enter_context(
+            mock.patch.object(
+                lockable_resources,
+                "try_reserve",
+                return_value=None,
+            )
+        )
+        mock_poll = exit_stack.enter_context(
+            mock.patch.object(
+                lockable_resources,
+                "poll",
+            )
+        )
+        exit_stack.enter_context(
+            pytest.raises(ResourceReservationTimeoutError)
+        )
+        with lockable_resources.reservation_by_name(
+            test_lock_name,
+            retry=SimpleRetryConfig(
+                sleep_period=1,
+                timeout=5.5,
+            ),
+        ):
+            pass
+    assert mock_time_sleep.call_count == 5
+    assert mock_try_reserve.call_count == 6
+    assert mock_poll.call_count == 5
