✨ Implement a dispenser locking mechanism

royalnet/engineer/command.py

@@ -32,7 +32,7 @@ class Command(c.Conversation):
               are later completed by a runner.
     """
 
-    def __init__(self, f: c.ConversationProtocol, *, names: t.List[str] = None, pattern: re.Pattern):
+    def __init__(self, f: c.ConversationProtocol, *, names: t.List[str] = None, pattern: re.Pattern, lock: bool = True):
         """
         Create a new :class:`.Command` .
         """
@@ -54,6 +54,11 @@ class Command(c.Conversation):
         The pattern that should be matched by the command.
         """
 
+        self.lock: bool = lock
+        """
+        If calling this command should :meth:`~royalnet.engineer.dispenser.Dispenser.lock` the dispenser.
+        """
+
     def name(self):
         """
         :return: The main name of the Command.
@@ -96,9 +101,17 @@ class Command(c.Conversation):
 
         log.debug(f"Match successful, getting capture groups of: {match!r}")
         message_kwargs: t.Dict[str, str] = match.groupdict()
-        
-        log.debug(f"Passing args to function: {message_kwargs!r}")
-        return await super().run(_sentry=_sentry, _msg=bullet, **base_kwargs, **message_kwargs)
+
+        if self.lock:
+            log.debug(f"Locking the dispenser...")
+
+            with _sentry.dispenser().lock(self):
+                log.debug(f"Passing args to function: {message_kwargs!r}")
+                return await super().run(_sentry=_sentry, _msg=bullet, **base_kwargs, **message_kwargs)
+
+        else:
+            log.debug(f"Passing args to function: {message_kwargs!r}")
+            return await super().run(_sentry=_sentry, _msg=bullet, **base_kwargs, **message_kwargs)
 
     def help(self) -> t.Optional[str]:
         """
@@ -116,7 +129,7 @@ class PartialCommand:
 
     They can specified later using :meth:`.complete`.
     """
-    def __init__(self, f: c.ConversationProtocol, syntax: str):
+    def __init__(self, f: c.ConversationProtocol, syntax: str, lock: bool = True):
         """
         Create a new :class:`.PartialCommand` .
 
@@ -133,6 +146,11 @@ class PartialCommand:
         Part of the pattern from where the arguments should be captured.
         """
 
+        self.lock: bool = lock
+        """
+        If calling this command should :meth:`~royalnet.engineer.dispenser.Dispenser.lock` the dispenser.
+        """
+
     @classmethod
     def new(cls, *args, **kwargs):
         """
@@ -167,8 +185,10 @@ class PartialCommand:
                 raise ValueError(f"Name is not alphanumeric: {name!r}")
 
         name_regex = f"(?:{'|'.join(names)})"
+        log.debug(f"Completed pattern: {name_regex!r}")
+
         pattern: re.Pattern = re.compile(pattern.format(name=name_regex, syntax=self.syntax), re.IGNORECASE)
-        return Command(f=self.f, names=names, pattern=pattern)
+        return Command(f=self.f, names=names, pattern=pattern, lock=self.lock)
 
     def __repr__(self):
         return f"<{self.__class__.__qualname__} {self.f!r}>"

royalnet/engineer/dispenser.py

@@ -10,6 +10,7 @@ import contextlib
 
 from .sentry import SentrySource
 from .conversation import Conversation
+from .exc import LockedDispenserError
 
 log = logging.getLogger(__name__)
 
@@ -21,6 +22,13 @@ class Dispenser:
         A :class:`list` of all the running sentries of this dispenser.
         """
 
+        self._locked_by: t.List[Conversation] = []
+        """
+        The conversation that are currently locking this dispenser.
+        
+        .. seealso:: :meth:`.lock`
+        """
+
     async def put(self, item: t.Any) -> None:
         """
         Insert a new item in the queues of all the running sentries.
@@ -54,7 +62,15 @@ class Dispenser:
         Run the passed conversation.
 
         :param conv: The conversation to run.
+        :raises .LockedDispenserError: If the dispenser is currently :attr:`.locked_by` a :class:`.Conversation`.
         """
+        log.debug(f"Trying to run: {conv!r}")
+
+        if self._locked_by is not None:
+            log.debug(f"Dispenser is locked by {self._locked_by!r}, refusing to run {conv!r}")
+            raise LockedDispenserError(self._locked_by, f"The Dispenser is currently locked by {self._locked_by!r} and "
+                                                        f"cannot start new conversations.")
+
         log.debug(f"Running: {conv}")
         with self.sentry() as sentry:
             state = conv(_sentry=sentry, **kwargs)
@@ -63,6 +79,26 @@ class Dispenser:
             while state := await state:
                 log.debug(f"Switched to: {state}")
 
+    @contextlib.contextmanager
+    def lock(self, conv: Conversation):
+        """
+        Lock the dispenser while this :func:`~contextlib.contextmanager` is in scope.
+
+        A locked dispenser will refuse to :meth:`.run` any new conversations, raising :exc:`.exc.LockedDispenserError`
+        instead.
+
+        :param conv: The conversation that requested the lock.
+
+        .. seealso:: :attr:`._locked_by`
+        """
+        log.debug(f"Adding lock: {conv!r}")
+        self._locked_by.append(conv)
+
+        yield
+
+        log.debug(f"Clearing lock: {conv!r}")
+        self._locked_by.remove(conv)
+
 
 __all__ = (
     "Dispenser",

royalnet/engineer/exc.py

@@ -61,6 +61,21 @@ class ForbiddenError(FrontendError):
     """
 
 
+class DispenserException(EngineerException):
+    """
+    The base class for errors in :mod:`royalnet.engineer.dispenser`.
+    """
+
+
+class LockedDispenserError(DispenserException):
+    """
+    The dispenser couldn't start a new conversation as it is currently locked.
+    """
+    def __init__(self, locked_by, *args):
+        super().__init__(*args)
+        self.locked_by = locked_by
+
+
 __all__ = (
     "EngineerException",
     "WrenchException",
@@ -72,4 +87,6 @@ __all__ = (
     "FrontendError",
     "NotSupportedError",
     "ForbiddenError",
+    "DispenserException",
+    "LockedDispenserError",
 )