rasa-pro 3.12.0.dev1__py3-none-any.whl

rasa/core/lock.py

@@ -0,0 +1,139 @@
+import json
+import logging
+from collections import deque
+
+import time
+from typing import Text, Optional, Union, Deque, Dict, Any
+
+logger = logging.getLogger(__name__)
+
+NO_TICKET_ISSUED = -1  # index of latest issued ticket if no tickets exist
+
+
+class Ticket:
+    def __init__(self, number: int, expires: float) -> None:
+        self.number = number
+        self.expires = expires
+
+    def has_expired(self) -> bool:
+        return time.time() > self.expires
+
+    def as_dict(self) -> Dict[Text, Any]:
+        return dict(number=self.number, expires=self.expires)
+
+    def dumps(self) -> Text:
+        """Return json dump of `Ticket` as dictionary."""
+        return json.dumps(self.as_dict())
+
+    @classmethod
+    def from_dict(cls, data: Dict[Text, Union[int, float]]) -> "Ticket":
+        """Creates `Ticket` from dictionary."""
+        return cls(number=data["number"], expires=data["expires"])
+
+    def __repr__(self) -> Text:
+        return f"Ticket(number: {self.number}, expires: {self.expires})"
+
+
+class TicketLock:
+    """Locking mechanism that issues tickets managing access to conversation IDs.
+
+    Tickets are issued in the order in which they are requested. A detailed
+    explanation of the ticket lock algorithm can be found at
+    http://pages.cs.wisc.edu/~remzi/OSTEP/threads-locks.pdf#page=13
+    """
+
+    def __init__(
+        self, conversation_id: Text, tickets: Optional[Deque[Ticket]] = None
+    ) -> None:
+        self.conversation_id = conversation_id
+        self.tickets = tickets or deque()
+
+    @classmethod
+    def from_dict(cls, data: Dict[Text, Any]) -> "TicketLock":
+        """Create `TicketLock` from dictionary."""
+        tickets = [Ticket.from_dict(json.loads(d)) for d in data.get("tickets", [])]
+        return cls(data.get("conversation_id"), deque(tickets))
+
+    def dumps(self) -> Text:
+        """Return json dump of `TicketLock`."""
+        tickets = [ticket.dumps() for ticket in self.tickets]
+        return json.dumps(dict(conversation_id=self.conversation_id, tickets=tickets))
+
+    def is_locked(self, ticket_number: int) -> bool:
+        """Return whether `ticket_number` is locked.
+
+        Returns:
+             True if `now_serving` is not equal to `ticket`.
+        """
+        return self.now_serving != ticket_number
+
+    def issue_ticket(self, lifetime: float) -> int:
+        """Issue a new ticket and return its number."""
+        self.remove_expired_tickets()
+        number = self.last_issued + 1
+        ticket = Ticket(number, time.time() + lifetime)
+        self.tickets.append(ticket)
+
+        return number
+
+    def remove_expired_tickets(self) -> None:
+        """Remove expired tickets."""
+        # iterate over copy of self.tickets so we can remove items
+        for ticket in list(self.tickets):
+            if ticket.has_expired():
+                self.tickets.remove(ticket)
+
+    @property
+    def last_issued(self) -> int:
+        """Return number of the ticket that was last added.
+
+        Returns:
+             Number of `Ticket` that was last added. `NO_TICKET_ISSUED` if no
+             tickets exist.
+        """
+        ticket_number = self._ticket_number_for(-1)
+
+        return ticket_number if ticket_number is not None else NO_TICKET_ISSUED
+
+    @property
+    def now_serving(self) -> Optional[int]:
+        """Get number of the ticket to be served next.
+
+        Returns:
+             Number of `Ticket` that is served next. 0 if no `Ticket` exists.
+        """
+        return self._ticket_number_for(0) or 0
+
+    def _ticket_number_for(self, ticket_index: int) -> Optional[int]:
+        """Get ticket number for `ticket_index`.
+
+        Returns:
+             Ticket number for `Ticket` with index `ticket_index`. None if there are no
+             tickets, or if `ticket_index` is out of bounds of `self.tickets`.
+        """
+        self.remove_expired_tickets()
+
+        try:
+            return self.tickets[ticket_index].number
+        except IndexError:
+            return None
+
+    def _ticket_for_ticket_number(self, ticket_number: int) -> Optional[Ticket]:
+        """Return ticket for `ticket_number`."""
+        self.remove_expired_tickets()
+
+        return next((t for t in self.tickets if t.number == ticket_number), None)
+
+    def is_someone_waiting(self) -> bool:
+        """Return whether someone is waiting for the lock to become available.
+
+        Returns:
+             True if the `self.tickets` queue has length greater than 0.
+        """
+        return len(self.tickets) > 0
+
+    def remove_ticket_for(self, ticket_number: int) -> None:
+        """Remove `Ticket` for `ticket_number."""
+        ticket = self._ticket_for_ticket_number(ticket_number)
+        if ticket:
+            self.tickets.remove(ticket)

rasa/core/lock_store.py

@@ -0,0 +1,343 @@
+from __future__ import annotations
+import asyncio
+from contextlib import asynccontextmanager
+import json
+import logging
+import os
+
+from typing import AsyncGenerator, Dict, Optional, Text, Union
+
+from rasa.shared.exceptions import RasaException, ConnectionException
+import rasa.shared.utils.common
+from rasa.core.constants import DEFAULT_LOCK_LIFETIME
+from rasa.core.lock import TicketLock
+from rasa.utils.endpoints import EndpointConfig
+
+logger = logging.getLogger(__name__)
+
+
+def _get_lock_lifetime() -> int:
+    return int(os.environ.get("TICKET_LOCK_LIFETIME", 0)) or DEFAULT_LOCK_LIFETIME
+
+
+LOCK_LIFETIME = _get_lock_lifetime()
+DEFAULT_SOCKET_TIMEOUT_IN_SECONDS = 10
+
+DEFAULT_REDIS_LOCK_STORE_KEY_PREFIX = "lock:"
+
+
+# noinspection PyUnresolvedReferences
+class LockError(RasaException):
+    """Exception that is raised when a lock cannot be acquired.
+
+    Attributes:
+         message (str): explanation of which `conversation_id` raised the error
+    """
+
+    pass
+
+
+class LockStore:
+    """Base class for ticket locks."""
+
+    @staticmethod
+    def create(obj: Union[LockStore, EndpointConfig, None]) -> LockStore:
+        """Factory to create a lock store."""
+        if isinstance(obj, LockStore):
+            return obj
+
+        try:
+            return _create_from_endpoint_config(obj)
+        except ConnectionError as error:
+            raise ConnectionException("Cannot connect to lock store.") from error
+
+    @staticmethod
+    def create_lock(conversation_id: Text) -> TicketLock:
+        """Create a new `TicketLock` for `conversation_id`."""
+        return TicketLock(conversation_id)
+
+    def get_lock(self, conversation_id: Text) -> Optional[TicketLock]:
+        """Fetch lock for `conversation_id` from storage."""
+        raise NotImplementedError
+
+    def delete_lock(self, conversation_id: Text) -> None:
+        """Delete lock for `conversation_id` from storage."""
+        raise NotImplementedError
+
+    def save_lock(self, lock: TicketLock) -> None:
+        """Commit `lock` to storage."""
+        raise NotImplementedError
+
+    def issue_ticket(
+        self, conversation_id: Text, lock_lifetime: float = LOCK_LIFETIME
+    ) -> int:
+        """Issue new ticket with `lock_lifetime` for lock associated with
+        `conversation_id`.
+
+        Creates a new lock if none is found.
+        """
+        logger.debug(f"Issuing ticket for conversation '{conversation_id}'.")
+        try:
+            lock = self.get_or_create_lock(conversation_id)
+            ticket = lock.issue_ticket(lock_lifetime)
+            self.save_lock(lock)
+
+            return ticket
+        except Exception as e:
+            raise LockError(f"Error while acquiring lock. Error:\n{e}")
+
+    @asynccontextmanager
+    async def lock(
+        self,
+        conversation_id: Text,
+        lock_lifetime: float = LOCK_LIFETIME,
+        wait_time_in_seconds: float = 1,
+    ) -> AsyncGenerator[TicketLock, None]:
+        """Acquire lock with lifetime `lock_lifetime`for `conversation_id`.
+
+        Try acquiring lock with a wait time of `wait_time_in_seconds` seconds
+        between attempts. Raise a `LockError` if lock has expired.
+        """
+        ticket = self.issue_ticket(conversation_id, lock_lifetime)
+        try:
+            yield await self._acquire_lock(
+                conversation_id, ticket, wait_time_in_seconds
+            )
+        finally:
+            self.cleanup(conversation_id, ticket)
+
+    async def _acquire_lock(
+        self, conversation_id: Text, ticket: int, wait_time_in_seconds: float
+    ) -> TicketLock:
+        logger.debug(f"Acquiring lock for conversation '{conversation_id}'.")
+        while True:
+            # fetch lock in every iteration because lock might no longer exist
+            lock = self.get_lock(conversation_id)
+
+            # exit loop if lock does not exist anymore (expired)
+            if not lock:
+                break
+
+            # acquire lock if it isn't locked
+            if not lock.is_locked(ticket):
+                logger.debug(f"Acquired lock for conversation '{conversation_id}'.")
+                return lock
+
+            items_before_this = ticket - (lock.now_serving or 0)
+
+            logger.debug(
+                f"Failed to acquire lock for conversation ID '{conversation_id}' "
+                f"because {items_before_this} other item(s) for this "
+                f"conversation ID have to be finished processing first. "
+                f"Retrying in {wait_time_in_seconds} seconds ..."
+            )
+
+            # sleep and update lock
+            await asyncio.sleep(wait_time_in_seconds)
+            self.update_lock(conversation_id)
+
+        raise LockError(
+            f"Could not acquire lock for conversation_id '{conversation_id}'."
+        )
+
+    def update_lock(self, conversation_id: Text) -> None:
+        """Fetch lock for `conversation_id`, remove expired tickets and save lock."""
+        lock = self.get_lock(conversation_id)
+        if lock:
+            lock.remove_expired_tickets()
+            self.save_lock(lock)
+
+    def get_or_create_lock(self, conversation_id: Text) -> TicketLock:
+        """Fetch existing lock for `conversation_id`.
+
+        Alternatively, create a new one if it doesn't exist.
+        """
+        existing_lock = self.get_lock(conversation_id)
+
+        if existing_lock:
+            return existing_lock
+
+        return self.create_lock(conversation_id)
+
+    def is_someone_waiting(self, conversation_id: Text) -> bool:
+        """Return whether someone is waiting for lock for this `conversation_id`."""
+        lock = self.get_lock(conversation_id)
+        if lock:
+            return lock.is_someone_waiting()
+
+        return False
+
+    def finish_serving(self, conversation_id: Text, ticket_number: int) -> None:
+        """Finish serving ticket with `ticket_number` for `conversation_id`.
+
+        Removes ticket from lock and saves lock.
+        """
+        lock = self.get_lock(conversation_id)
+        if lock:
+            lock.remove_ticket_for(ticket_number)
+            self.save_lock(lock)
+
+    def cleanup(self, conversation_id: Text, ticket_number: int) -> None:
+        """Remove lock for `conversation_id` if no one is waiting."""
+        self.finish_serving(conversation_id, ticket_number)
+        if not self.is_someone_waiting(conversation_id):
+            self.delete_lock(conversation_id)
+
+    @staticmethod
+    def _log_deletion(conversation_id: Text, deletion_successful: bool) -> None:
+        if deletion_successful:
+            logger.debug(f"Deleted lock for conversation '{conversation_id}'.")
+        else:
+            logger.debug(f"Could not delete lock for conversation '{conversation_id}'.")
+
+
+class RedisLockStore(LockStore):
+    """Redis store for ticket locks."""
+
+    def __init__(
+        self,
+        host: Text = "localhost",
+        port: int = 6379,
+        db: int = 1,
+        username: Optional[Text] = None,
+        password: Optional[Text] = None,
+        use_ssl: bool = False,
+        ssl_certfile: Optional[Text] = None,
+        ssl_keyfile: Optional[Text] = None,
+        ssl_ca_certs: Optional[Text] = None,
+        key_prefix: Optional[Text] = None,
+        socket_timeout: float = DEFAULT_SOCKET_TIMEOUT_IN_SECONDS,
+    ) -> None:
+        """Create a lock store which uses Redis for persistence.
+
+        Args:
+            host: The host of the redis server.
+            port: The port of the redis server.
+            db: The name of the database within Redis which should be used by Rasa
+                Open Source.
+            username: The username which should be used for authentication with the
+                Redis database.
+            password: The password which should be used for authentication with the
+                Redis database.
+            use_ssl: `True` if SSL should be used for the connection to Redis.
+            ssl_certfile: Path to the SSL certificate file.
+            ssl_keyfile: Path to the SSL private key file.
+            ssl_ca_certs: Path to the SSL CA certificate file.
+            key_prefix: prefix to prepend to all keys used by the lock store. Must be
+                alphanumeric.
+            socket_timeout: Timeout in seconds after which an exception will be raised
+                in case Redis doesn't respond within `socket_timeout` seconds.
+        """
+        import redis
+
+        self.red = redis.StrictRedis(
+            host=host,
+            port=int(port),
+            db=int(db),
+            username=username,
+            password=password,
+            ssl=use_ssl,
+            ssl_certfile=ssl_certfile,
+            ssl_keyfile=ssl_keyfile,
+            ssl_ca_certs=ssl_ca_certs,
+            socket_timeout=socket_timeout,
+        )
+
+        self.key_prefix = DEFAULT_REDIS_LOCK_STORE_KEY_PREFIX
+        if key_prefix:
+            logger.debug(f"Setting non-default redis key prefix: '{key_prefix}'.")
+            self._set_key_prefix(key_prefix)
+
+        super().__init__()
+
+    def _set_key_prefix(self, key_prefix: Text) -> None:
+        if isinstance(key_prefix, str) and key_prefix.isalnum():
+            self.key_prefix = key_prefix + ":" + DEFAULT_REDIS_LOCK_STORE_KEY_PREFIX
+        else:
+            logger.warning(
+                f"Omitting provided non-alphanumeric redis key prefix: '{key_prefix}'. "
+                f"Using default '{self.key_prefix}' instead."
+            )
+
+    def get_lock(self, conversation_id: Text) -> Optional[TicketLock]:
+        """Retrieves lock (see parent docstring for more information)."""
+        serialised_lock = self.red.get(self.key_prefix + conversation_id)
+        if serialised_lock:
+            return TicketLock.from_dict(json.loads(serialised_lock))
+
+        return None
+
+    def delete_lock(self, conversation_id: Text) -> None:
+        """Deletes lock for conversation ID."""
+        deletion_successful = self.red.delete(self.key_prefix + conversation_id)
+        self._log_deletion(conversation_id, deletion_successful)
+
+    def save_lock(self, lock: TicketLock) -> None:
+        self.red.set(self.key_prefix + lock.conversation_id, lock.dumps())
+
+
+class InMemoryLockStore(LockStore):
+    """In-memory store for ticket locks."""
+
+    def __init__(self) -> None:
+        """Initialise dictionary of locks."""
+        self.conversation_locks: Dict[Text, TicketLock] = {}
+        super().__init__()
+
+    def get_lock(self, conversation_id: Text) -> Optional[TicketLock]:
+        """Get lock for conversation if it exists."""
+        return self.conversation_locks.get(conversation_id)
+
+    def delete_lock(self, conversation_id: Text) -> None:
+        """Delete lock for conversation."""
+        deleted_lock = self.conversation_locks.pop(conversation_id, None)
+        self._log_deletion(
+            conversation_id, deletion_successful=deleted_lock is not None
+        )
+
+    def save_lock(self, lock: TicketLock) -> None:
+        """Save lock in store."""
+        self.conversation_locks[lock.conversation_id] = lock
+
+
+def _create_from_endpoint_config(
+    endpoint_config: Optional[EndpointConfig] = None,
+) -> LockStore:
+    """Given an endpoint configuration, create a proper `LockStore` object."""
+    if (
+        endpoint_config is None
+        or endpoint_config.type is None
+        or endpoint_config.type == "in_memory"
+    ):
+        # this is the default type if no lock store type is set
+
+        lock_store: LockStore = InMemoryLockStore()
+    elif endpoint_config.type == "redis":
+        lock_store = RedisLockStore(host=endpoint_config.url, **endpoint_config.kwargs)
+    elif endpoint_config.type == "concurrent_redis":
+        from rasa.core.concurrent_lock_store import ConcurrentRedisLockStore
+
+        lock_store = ConcurrentRedisLockStore(endpoint_config=endpoint_config)
+    else:
+        lock_store = _load_from_module_name_in_endpoint_config(endpoint_config)
+
+    logger.debug(f"Connected to lock store '{lock_store.__class__.__name__}'.")
+
+    return lock_store
+
+
+def _load_from_module_name_in_endpoint_config(
+    endpoint_config: EndpointConfig,
+) -> LockStore:
+    """Retrieve a `LockStore` based on its class name."""
+    try:
+        lock_store_class = rasa.shared.utils.common.class_from_module_path(
+            endpoint_config.type
+        )
+        return lock_store_class(endpoint_config=endpoint_config)
+    except (AttributeError, ImportError) as e:
+        raise Exception(
+            f"Could not find a class based on the module path "
+            f"'{endpoint_config.type}'. Failed to create a `LockStore` "
+            f"instance. Error: {e}"
+        )