nucliadb 6.4.0.post4200__py3-none-any.whl → 6.4.0.post4204__py3-none-any.whl

nucliadb/common/back_pressure/__init__.py

@@ -0,0 +1,20 @@
+# Copyright (C) 2021 Bosutech XXI S.L.
+#
+# nucliadb is offered under the AGPL v3.0 and as commercial software.
+# For commercial licensing, contact us at info@nuclia.com.
+#
+# AGPL:
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program 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 Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program. If not, see <http://www.gnu.org/licenses/>.
+#
+from .materializer import maybe_back_pressure, start_materializer, stop_materializer  # noqa

nucliadb/common/back_pressure/cache.py

@@ -0,0 +1,86 @@
+# Copyright (C) 2021 Bosutech XXI S.L.
+#
+# nucliadb is offered under the AGPL v3.0 and as commercial software.
+# For commercial licensing, contact us at info@nuclia.com.
+#
+# AGPL:
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program 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 Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program. If not, see <http://www.gnu.org/licenses/>.
+#
+import contextlib
+import logging
+import threading
+from datetime import datetime, timezone
+from typing import Optional
+
+from cachetools import TTLCache
+
+from nucliadb.common.back_pressure.utils import BackPressureData, BackPressureException
+from nucliadb_telemetry import metrics
+
+logger = logging.getLogger(__name__)
+
+
+RATE_LIMITED_REQUESTS_COUNTER = metrics.Counter(
+    "nucliadb_rate_limited_requests", labels={"type": "", "cached": ""}
+)
+
+
+class BackPressureCache:
+    """
+    Global cache for storing already computed try again in times.
+    It allows us to avoid making the same calculations multiple
+    times if back pressure has been applied.
+    """
+
+    def __init__(self):
+        self._cache = TTLCache(maxsize=1024, ttl=5 * 60)
+        self._lock = threading.Lock()
+
+    def get(self, key: str) -> Optional[BackPressureData]:
+        with self._lock:
+            data = self._cache.get(key, None)
+            if data is None:
+                return None
+            if datetime.now(timezone.utc) >= data.try_after:
+                # The key has expired, so remove it from the cache
+                self._cache.pop(key, None)
+                return None
+            return data
+
+    def set(self, key: str, data: BackPressureData):
+        with self._lock:
+            self._cache[key] = data
+
+
+_cache = BackPressureCache()
+
+
+@contextlib.contextmanager
+def cached_back_pressure(cache_key: str):
+    """
+    Context manager that handles the caching of the try again in time so that
+    we don't recompute try again times if we have already applied back pressure.
+    """
+    data: Optional[BackPressureData] = _cache.get(cache_key)
+    if data is not None:
+        back_pressure_type = data.type
+        RATE_LIMITED_REQUESTS_COUNTER.inc({"type": back_pressure_type, "cached": "true"})
+        raise BackPressureException(data)
+    try:
+        yield
+    except BackPressureException as exc:
+        back_pressure_type = exc.data.type
+        RATE_LIMITED_REQUESTS_COUNTER.inc({"type": back_pressure_type, "cached": "false"})
+        _cache.set(cache_key, exc.data)
+        raise exc

nucliadb/common/back_pressure/materializer.py

@@ -0,0 +1,315 @@
+# Copyright (C) 2021 Bosutech XXI S.L.
+#
+# nucliadb is offered under the AGPL v3.0 and as commercial software.
+# For commercial licensing, contact us at info@nuclia.com.
+#
+# AGPL:
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program 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 Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program. If not, see <http://www.gnu.org/licenses/>.
+#
+import asyncio
+import logging
+import threading
+from typing import Optional
+
+from cachetools import TTLCache
+from fastapi import HTTPException
+
+from nucliadb.common.back_pressure.cache import cached_back_pressure
+from nucliadb.common.back_pressure.settings import settings
+from nucliadb.common.back_pressure.utils import (
+    BackPressureData,
+    BackPressureException,
+    estimate_try_after,
+    get_nats_consumer_pending_messages,
+    is_back_pressure_enabled,
+)
+from nucliadb.common.context import ApplicationContext
+from nucliadb.common.http_clients.processing import ProcessingHTTPClient
+from nucliadb_telemetry import metrics
+from nucliadb_utils import const
+from nucliadb_utils.nats import NatsConnectionManager
+from nucliadb_utils.settings import is_onprem_nucliadb
+
+logger = logging.getLogger(__name__)
+
+
+back_pressure_observer = metrics.Observer("nucliadb_back_pressure", labels={"type": ""})
+
+
+class BackPressureMaterializer:
+    """
+    Singleton class that will run in the background gathering the different
+    stats to apply back pressure and materializing it in memory. This allows us
+    to do stale-reads when checking if back pressure is needed for a particular
+    request - thus not slowing it down.
+    """
+
+    def __init__(
+        self,
+        nats_manager: NatsConnectionManager,
+        indexing_check_interval: int = 30,
+        ingest_check_interval: int = 30,
+    ):
+        self.nats_manager = nats_manager
+        self.processing_http_client = ProcessingHTTPClient()
+
+        self.indexing_check_interval = indexing_check_interval
+        self.ingest_check_interval = ingest_check_interval
+
+        self.ingest_pending: int = 0
+        self.indexing_pending: int = 0
+
+        self._tasks: list[asyncio.Task] = []
+        self._running = False
+
+        self.processing_pending_cache = TTLCache(maxsize=1024, ttl=60)  # type: ignore
+        self.processing_pending_locks: dict[str, asyncio.Lock] = {}
+
+    async def start(self):
+        self._tasks.append(asyncio.create_task(self._get_indexing_pending_task()))
+        self._tasks.append(asyncio.create_task(self._get_ingest_pending_task()))
+        self._running = True
+
+    async def stop(self):
+        for task in self._tasks:
+            task.cancel()
+        self._tasks.clear()
+        await self.processing_http_client.close()
+        self._running = False
+
+    @property
+    def running(self) -> bool:
+        return self._running
+
+    async def get_processing_pending(self, kbid: str) -> int:
+        """
+        We don't materialize the pending messages for every kbid, but values are cached for some time.
+        """
+        cached = self.processing_pending_cache.get(kbid)
+        if cached is not None:
+            return cached
+
+        lock = self.processing_pending_locks.setdefault(kbid, asyncio.Lock())
+        async with lock:
+            # Check again if the value has been cached while we were waiting for the lock
+            cached = self.processing_pending_cache.get(kbid)
+            if cached is not None:
+                return cached
+
+            # Get the pending messages and cache the result
+            try:
+                with back_pressure_observer({"type": "get_processing_pending"}):
+                    pending = await self._get_processing_pending(kbid)
+            except Exception:  # pragma: no cover
+                # Do not cache if there was an error
+                logger.exception(
+                    "Error getting pending messages to process. Back pressure on proccessing for KB can't be applied.",
+                    exc_info=True,
+                    extra={"kbid": kbid},
+                )
+                return 0
+
+            if pending > 0:
+                logger.info(
+                    f"Processing returned {pending} pending messages for KB",
+                    extra={"kbid": kbid},
+                )
+            self.processing_pending_cache[kbid] = pending
+            return pending
+
+    async def _get_processing_pending(self, kbid: str) -> int:
+        response = await self.processing_http_client.stats(kbid=kbid, timeout=0.5)
+        return response.incomplete
+
+    def get_indexing_pending(self) -> int:
+        return self.indexing_pending
+
+    def get_ingest_pending(self) -> int:
+        return self.ingest_pending
+
+    async def _get_indexing_pending_task(self):
+        try:
+            while True:
+                try:
+                    with back_pressure_observer({"type": "get_indexing_pending"}):
+                        self.indexing_pending = await get_nats_consumer_pending_messages(
+                            self.nats_manager,
+                            stream="nidx",
+                            consumer="nidx",
+                        )
+                except Exception:  # pragma: no cover
+                    logger.exception(
+                        "Error getting pending messages to index",
+                        exc_info=True,
+                    )
+                await asyncio.sleep(self.indexing_check_interval)
+        except asyncio.CancelledError:
+            pass
+
+    async def _get_ingest_pending_task(self):
+        try:
+            while True:
+                try:
+                    with back_pressure_observer({"type": "get_ingest_pending"}):
+                        self.ingest_pending = await get_nats_consumer_pending_messages(
+                            self.nats_manager,
+                            stream=const.Streams.INGEST_PROCESSED.name,
+                            consumer=const.Streams.INGEST_PROCESSED.group,
+                        )
+                except Exception:  # pragma: no cover
+                    logger.exception(
+                        "Error getting pending messages to ingest",
+                        exc_info=True,
+                    )
+                await asyncio.sleep(self.ingest_check_interval)
+        except asyncio.CancelledError:
+            pass
+
+    def check_indexing(self):
+        max_pending = settings.max_indexing_pending
+        if max_pending <= 0:
+            # Indexing back pressure is disabled
+            return
+        pending = self.get_indexing_pending()
+        if pending > max_pending:
+            try_after = estimate_try_after(
+                rate=settings.indexing_rate,
+                pending=pending,
+                max_wait=settings.max_wait_time,
+            )
+            data = BackPressureData(type="indexing", try_after=try_after)
+            raise BackPressureException(data)
+
+    def check_ingest(self):
+        max_pending = settings.max_ingest_pending
+        if max_pending <= 0:
+            # Ingest back pressure is disabled
+            return
+        ingest_pending = self.get_ingest_pending()
+        if ingest_pending > max_pending:
+            try_after = estimate_try_after(
+                rate=settings.ingest_rate,
+                pending=ingest_pending,
+                max_wait=settings.max_wait_time,
+            )
+            data = BackPressureData(type="ingest", try_after=try_after)
+            raise BackPressureException(data)
+
+    async def check_processing(self, kbid: str):
+        max_pending = settings.max_processing_pending
+        if max_pending <= 0:
+            # Processing back pressure is disabled
+            return
+
+        kb_pending = await self.get_processing_pending(kbid)
+        if kb_pending > max_pending:
+            try_after = estimate_try_after(
+                rate=settings.processing_rate,
+                pending=kb_pending,
+                max_wait=settings.max_wait_time,
+            )
+            data = BackPressureData(type="processing", try_after=try_after)
+            raise BackPressureException(data)
+
+
+MATERIALIZER: Optional[BackPressureMaterializer] = None
+materializer_lock = threading.Lock()
+
+
+async def start_materializer(context: ApplicationContext):
+    global MATERIALIZER
+    if MATERIALIZER is not None:
+        logger.warning("BackPressureMaterializer already started")
+        return
+    with materializer_lock:
+        if MATERIALIZER is not None:
+            return
+        logger.info("Initializing materializer")
+        try:
+            nats_manager = context.nats_manager
+        except AttributeError:
+            logger.warning(
+                "Could not initialize materializer. Nats manager not found or not initialized yet"
+            )
+            return
+        materializer = BackPressureMaterializer(
+            nats_manager,
+            indexing_check_interval=settings.indexing_check_interval,
+            ingest_check_interval=settings.ingest_check_interval,
+        )
+        await materializer.start()
+        MATERIALIZER = materializer
+
+
+async def stop_materializer():
+    global MATERIALIZER
+    if MATERIALIZER is None or not MATERIALIZER.running:
+        logger.warning("BackPressureMaterializer already stopped")
+        return
+    with materializer_lock:
+        if MATERIALIZER is None:
+            return
+        logger.info("Stopping materializer")
+        await MATERIALIZER.stop()
+        MATERIALIZER = None
+
+
+def get_materializer() -> BackPressureMaterializer:
+    global MATERIALIZER
+    if MATERIALIZER is None:
+        raise RuntimeError("BackPressureMaterializer not initialized")
+    return MATERIALIZER
+
+
+async def maybe_back_pressure(kbid: str, resource_uuid: Optional[str] = None) -> None:
+    """
+    This function does system checks to see if we need to put back pressure on writes.
+    In that case, a HTTP 429 will be raised with the estimated time to try again.
+    """
+    if not is_back_pressure_enabled() or is_onprem_nucliadb():
+        return
+    await back_pressure_checks(kbid, resource_uuid)
+
+
+async def back_pressure_checks(kbid: str, resource_uuid: Optional[str] = None):
+    """
+    Will raise a 429 if back pressure is needed:
+    - If the processing engine is behind.
+    - If ingest processed consumer is behind.
+    - If the indexing on nodes affected by the request (kbid, and resource_uuid) is behind.
+    """
+    materializer = get_materializer()
+    try:
+        with cached_back_pressure(f"{kbid}-{resource_uuid}"):
+            materializer.check_indexing()
+            materializer.check_ingest()
+            await materializer.check_processing(kbid)
+    except BackPressureException as exc:
+        logger.info(
+            "Back pressure applied",
+            extra={
+                "kbid": kbid,
+                "resource_uuid": resource_uuid,
+                "try_after": exc.data.try_after,
+                "back_pressure_type": exc.data.type,
+            },
+        )
+        raise HTTPException(
+            status_code=429,
+            detail={
+                "message": f"Too many messages pending to ingest. Retry after {exc.data.try_after}",
+                "try_after": exc.data.try_after.timestamp(),
+                "back_pressure_type": exc.data.type,
+            },
+        ) from exc

nucliadb/common/back_pressure/settings.py

@@ -0,0 +1,72 @@
+# Copyright (C) 2021 Bosutech XXI S.L.
+#
+# nucliadb is offered under the AGPL v3.0 and as commercial software.
+# For commercial licensing, contact us at info@nuclia.com.
+#
+# AGPL:
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program 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 Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program. If not, see <http://www.gnu.org/licenses/>.
+#
+from pydantic import Field
+from pydantic_settings import BaseSettings
+
+
+class BackPressureSettings(BaseSettings):
+    enabled: bool = Field(
+        default=False,
+        description="Enable or disable back pressure.",
+        alias="back_pressure_enabled",
+    )
+    indexing_rate: float = Field(
+        default=10,
+        description="Estimation of the indexing rate in messages per second. This is used to calculate the try again in time",  # noqa
+    )
+    ingest_rate: float = Field(
+        default=4,
+        description="Estimation of the ingest processed consumer rate in messages per second. This is used to calculate the try again in time",  # noqa
+    )
+    processing_rate: float = Field(
+        default=1,
+        description="Estimation of the processing rate in messages per second. This is used to calculate the try again in time",  # noqa
+    )
+    max_indexing_pending: int = Field(
+        default=1000,
+        description="Max number of messages pending to index in a node queue before rate limiting writes. Set to 0 to disable indexing back pressure checks",  # noqa
+        alias="back_pressure_max_indexing_pending",
+    )
+    max_ingest_pending: int = Field(
+        # Disabled by default
+        default=0,
+        description="Max number of messages pending to be ingested by processed consumers before rate limiting writes. Set to 0 to disable ingest back pressure checks",  # noqa
+        alias="back_pressure_max_ingest_pending",
+    )
+    max_processing_pending: int = Field(
+        default=1000,
+        description="Max number of messages pending to process per Knowledge Box before rate limiting writes. Set to 0 to disable processing back pressure checks",  # noqa
+        alias="back_pressure_max_processing_pending",
+    )
+    indexing_check_interval: int = Field(
+        default=30,
+        description="Interval in seconds to check the indexing pending messages",
+    )
+    ingest_check_interval: int = Field(
+        default=30,
+        description="Interval in seconds to check the ingest pending messages",
+    )
+    max_wait_time: int = Field(
+        default=60,
+        description="Max time in seconds to wait before trying again after back pressure",
+    )
+
+
+settings = BackPressureSettings()

nucliadb/common/back_pressure/utils.py

@@ -0,0 +1,59 @@
+# Copyright (C) 2021 Bosutech XXI S.L.
+#
+# nucliadb is offered under the AGPL v3.0 and as commercial software.
+# For commercial licensing, contact us at info@nuclia.com.
+#
+# AGPL:
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program 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 Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program. If not, see <http://www.gnu.org/licenses/>.
+#
+from dataclasses import dataclass
+from datetime import datetime, timedelta, timezone
+
+from nucliadb.common.back_pressure.settings import settings
+from nucliadb_utils.nats import NatsConnectionManager
+
+
+@dataclass
+class BackPressureData:
+    type: str
+    try_after: datetime
+
+
+class BackPressureException(Exception):
+    def __init__(self, data: BackPressureData):
+        self.data = data
+
+    def __str__(self):
+        return f"Back pressure applied for {self.data.type}. Try again after {self.data.try_after}"
+
+
+def is_back_pressure_enabled() -> bool:
+    return settings.enabled
+
+
+def estimate_try_after(rate: float, pending: int, max_wait: int) -> datetime:
+    """
+    This function estimates the time to try again based on the rate and the number of pending messages.
+    """
+    delta_seconds = min(pending / rate, max_wait)
+    return datetime.now(timezone.utc) + timedelta(seconds=delta_seconds)
+
+
+async def get_nats_consumer_pending_messages(
+    nats_manager: NatsConnectionManager, *, stream: str, consumer: str
+) -> int:
+    # get raw js client
+    js = nats_manager.js
+    consumer_info = await js.consumer_info(stream, consumer)
+    return consumer_info.num_pending

nucliadb/writer/api/v1/export_import.py

@@ -25,6 +25,7 @@ from fastapi_versioning import version
 from starlette.requests import Request
 
 from nucliadb.common import datamanagers
+from nucliadb.common.back_pressure import maybe_back_pressure
 from nucliadb.common.cluster.settings import in_standalone_mode
 from nucliadb.common.context import ApplicationContext
 from nucliadb.common.context.fastapi import get_app_context
@@ -45,7 +46,6 @@ from nucliadb.writer import logger
 from nucliadb.writer.api.utils import only_for_onprem
 from nucliadb.writer.api.v1.knowledgebox import create_kb
 from nucliadb.writer.api.v1.router import KB_PREFIX, KBS_PREFIX, api
-from nucliadb.writer.back_pressure import maybe_back_pressure
 from nucliadb_models.export_import import (
     CreateExportResponse,
     CreateImportResponse,
@@ -148,7 +148,7 @@ async def start_kb_import_endpoint(request: Request, kbid: str):
     if not await datamanagers.atomic.kb.exists_kb(kbid=kbid):
         return HTTPClientError(status_code=404, detail="Knowledge Box not found")
 
-    await maybe_back_pressure(request, kbid)
+    await maybe_back_pressure(kbid)
 
     stream = stream_compatible_with_kb(kbid, request.stream())
     try:

nucliadb/writer/api/v1/field.py

@@ -25,6 +25,7 @@ from fastapi_versioning import version
 from starlette.requests import Request
 
 import nucliadb_models as models
+from nucliadb.common.back_pressure import maybe_back_pressure
 from nucliadb.common.maindb.utils import get_driver
 from nucliadb.ingest.orm.knowledgebox import KnowledgeBox
 from nucliadb.models.internal.processing import PushPayload, Source
@@ -40,7 +41,6 @@ from nucliadb.writer.api.v1.resource import (
     validate_rid_exists_or_raise_error,
 )
 from nucliadb.writer.api.v1.router import KB_PREFIX, RESOURCE_PREFIX, RSLUG_PREFIX, api
-from nucliadb.writer.back_pressure import maybe_back_pressure
 from nucliadb.writer.resource.audit import parse_audit
 from nucliadb.writer.resource.field import (
     ResourceClassifications,
@@ -96,7 +96,7 @@ async def add_field_to_resource(
     **parser_kwargs,
 ):
     await validate_rid_exists_or_raise_error(kbid, rid)
-    await maybe_back_pressure(request, kbid, resource_uuid=rid)
+    await maybe_back_pressure(kbid, resource_uuid=rid)
 
     partitioning = get_partitioning()
     partition = partitioning.generate_partition(kbid, rid)
@@ -542,7 +542,7 @@ async def reprocess_file_field(
     x_nucliadb_user: Annotated[str, X_NUCLIADB_USER] = "",
     x_file_password: Annotated[Optional[str], X_FILE_PASSWORD] = None,
 ) -> ResourceUpdated:
-    await maybe_back_pressure(request, kbid, resource_uuid=rid)
+    await maybe_back_pressure(kbid, resource_uuid=rid)
 
     processing = get_processing()
     partitioning = get_partitioning()

nucliadb/writer/api/v1/resource.py

@@ -28,6 +28,7 @@ from fastapi_versioning import version
 from starlette.requests import Request
 
 from nucliadb.common import datamanagers
+from nucliadb.common.back_pressure import maybe_back_pressure
 from nucliadb.common.context.fastapi import get_app_context
 from nucliadb.common.maindb.driver import Driver
 from nucliadb.common.maindb.exceptions import ConflictError, NotFoundError
@@ -45,7 +46,6 @@ from nucliadb.writer.api.v1.router import (
     api,
 )
 from nucliadb.writer.api.v1.slug import ensure_slug_uniqueness, noop_context_manager
-from nucliadb.writer.back_pressure import maybe_back_pressure
 from nucliadb.writer.resource.audit import parse_audit
 from nucliadb.writer.resource.basic import (
     parse_basic_creation,
@@ -106,7 +106,7 @@ async def create_resource(
             detail="Cannot hide a resource: the KB does not have hidden resources enabled",
         )
 
-    await maybe_back_pressure(request, kbid)
+    await maybe_back_pressure(kbid)
 
     partitioning = get_partitioning()
 
@@ -243,7 +243,7 @@ async def modify_resource_endpoint(
 ):
     await validate_rid_exists_or_raise_error(kbid, rid)
 
-    await maybe_back_pressure(request, kbid, resource_uuid=rid)
+    await maybe_back_pressure(kbid, resource_uuid=rid)
 
     if item.slug is None:
         return await modify_resource(
@@ -424,7 +424,7 @@ async def _reprocess_resource(
     x_nucliadb_user: str,
 ):
     await validate_rid_exists_or_raise_error(kbid, rid)
-    await maybe_back_pressure(request, kbid, resource_uuid=rid)
+    await maybe_back_pressure(kbid, resource_uuid=rid)
 
     partitioning = get_partitioning()
 
@@ -571,7 +571,7 @@ async def _reindex_resource(
     reindex_vectors: bool,
 ):
     await validate_rid_exists_or_raise_error(kbid, rid)
-    await maybe_back_pressure(request, kbid, resource_uuid=rid)
+    await maybe_back_pressure(kbid, resource_uuid=rid)
 
     ingest = get_ingest()
     index_req = IndexResource()