nucliadb 6.9.1.post5192__py3-none-any.whl → 6.9.1.post5208__py3-none-any.whl

nucliadb/common/cluster/rebalance.py

@@ -18,165 +18,494 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 import asyncio
+import dataclasses
 import logging
+import math
+import random
+from typing import Optional
 
+import aioitertools
+from grpc import StatusCode
+from grpc.aio import AioRpcError
 from nidx_protos import nodereader_pb2, noderesources_pb2
 
 from nucliadb.common import datamanagers, locking
 from nucliadb.common.cluster.utils import get_shard_manager
 from nucliadb.common.context import ApplicationContext
+from nucliadb.common.datamanagers.resources import KB_RESOURCE_SHARD
 from nucliadb.common.nidx import get_nidx_api_client, get_nidx_searcher_client
+from nucliadb_protos import writer_pb2
 from nucliadb_telemetry import errors
 from nucliadb_telemetry.logs import setup_logging
 from nucliadb_telemetry.utils import setup_telemetry
+from nucliadb_utils import const
 from nucliadb_utils.fastapi.run import serve_metrics
+from nucliadb_utils.utilities import has_feature
 
 from .settings import settings
-from .utils import delete_resource_from_shard, index_resource_to_shard
+from .utils import delete_resource_from_shard, index_resource_to_shard, wait_for_nidx
 
 logger = logging.getLogger(__name__)
 
 REBALANCE_LOCK = "rebalance"
 
-
-async def get_shards_paragraphs(kbid: str) -> list[tuple[str, int]]:
-    """
-    Ordered shard -> num paragraph by number of paragraphs
-    """
-    async with datamanagers.with_ro_transaction() as txn:
-        kb_shards = await datamanagers.cluster.get_kb_shards(txn, kbid=kbid)
-    if kb_shards is None:
-        return []
-
-    results = {}
-    for shard_meta in kb_shards.shards:
-        # Rebalance using node as source of truth. But it will rebalance nidx
-        shard_data: nodereader_pb2.Shard = await get_nidx_api_client().GetShard(
-            nodereader_pb2.GetShardRequest(
-                shard_id=noderesources_pb2.ShardId(id=shard_meta.nidx_shard_id)
-            )  # type: ignore
+MAX_MOVES_PER_SHARD = 100
+
+
+@dataclasses.dataclass
+class RebalanceShard:
+    id: str
+    nidx_id: str
+    paragraphs: int
+    active: bool
+
+    def to_dict(self):
+        return self.__dict__
+
+
+class Rebalancer:
+    def __init__(self, context: ApplicationContext, kbid: str):
+        self.context = context
+        self.kbid = kbid
+        self.kb_shards: Optional[writer_pb2.Shards] = None
+        self.index: dict[str, set[str]] = {}
+
+    async def get_rebalance_shards(self) -> list[RebalanceShard]:
+        """
+        Return the sorted list of shards by increasing paragraph count.
+        """
+        self.kb_shards = await datamanagers.atomic.cluster.get_kb_shards(kbid=self.kbid)
+        if self.kb_shards is None:  # pragma: no cover
+            return []
+        return list(
+            sorted(
+                [
+                    RebalanceShard(
+                        id=shard.shard,
+                        nidx_id=shard.nidx_shard_id,
+                        paragraphs=await get_shard_paragraph_count(shard.nidx_shard_id),
+                        active=(idx == self.kb_shards.actual),
+                    )
+                    for idx, shard in enumerate(self.kb_shards.shards)
+                ],
+                key=lambda x: x.paragraphs,
+            )
         )
-        results[shard_meta.shard] = shard_data.paragraphs
 
-    return [(shard, paragraphs) for shard, paragraphs in sorted(results.items(), key=lambda x: x[1])]
+    async def build_shard_resources_index(self):
+        async with datamanagers.with_ro_transaction() as txn:
+            iterable = datamanagers.resources.iterate_resource_ids(kbid=self.kbid)
+            async for resources_batch in aioitertools.batched(iterable, n=200):
+                shards = await txn.batch_get(
+                    keys=[KB_RESOURCE_SHARD.format(kbid=self.kbid, uuid=rid) for rid in resources_batch],
+                    for_update=False,
+                )
+                for rid, shard_bytes in zip(resources_batch, shards):
+                    if shard_bytes is not None:
+                        self.index.setdefault(shard_bytes.decode(), set()).add(rid)
+
+    async def move_paragraphs(
+        self, from_shard: RebalanceShard, to_shard: RebalanceShard, max_paragraphs: int
+    ) -> int:
+        """
+        Takes random resources from the source shard and tries to move at most max_paragraphs.
+        It stops moving paragraphs until the are no more resources to move.
+        """
+        moved_paragraphs = 0
+
+        while moved_paragraphs < max_paragraphs:
+            # Take a random resource to move
+            try:
+                resource_id = random.choice(tuple(self.index[from_shard.id]))
+            except (KeyError, IndexError):
+                # No more resources in shard or shard not found
+                break
+
+            assert self.kb_shards is not None
+            from_shard_obj = next(s for s in self.kb_shards.shards if s.shard == from_shard.id)
+            to_shard_obj = next(s for s in self.kb_shards.shards if s.shard == to_shard.id)
+            paragraphs_count = await get_resource_paragraphs_count(resource_id, from_shard.nidx_id)
+            moved = await move_resource_to_shard(
+                self.context, self.kbid, resource_id, from_shard_obj, to_shard_obj
+            )
+            if moved:
+                self.index[from_shard.id].remove(resource_id)
+                self.index.setdefault(to_shard.id, set()).add(resource_id)
+                moved_paragraphs += paragraphs_count
 
+        return moved_paragraphs
 
-async def maybe_add_shard(kbid: str) -> None:
-    async with locking.distributed_lock(locking.NEW_SHARD_LOCK.format(kbid=kbid)):
-        async with datamanagers.with_ro_transaction() as txn:
-            kb_shards = await datamanagers.cluster.get_kb_shards(txn, kbid=kbid)
-        if kb_shards is None:
+    async def wait_for_indexing(self):
+        try:
+            self.context.nats_manager
+        except AssertionError:  # pragma: no cover
+            logger.warning(f"Nats manager not initialized. Cannot wait for indexing")
             return
+        while True:
+            try:
+                await wait_for_nidx(self.context.nats_manager, max_wait_seconds=60, max_pending=1000)
+                return
+            except asyncio.TimeoutError:
+                logger.warning("Nidx is behind. Backing off rebalancing.", extra={"kbid": self.kbid})
+                await asyncio.sleep(30)
+
+    async def rebalance_shards(self):
+        """
+        Iterate over shards until none of them need more rebalancing.
+
+        Will move excess of paragraphs to other shards (potentially creating new ones), and
+        merge small shards together when possible (potentially deleting empty ones.)
+
+
+        Merge chooses a <90% filled shard and fills it to almost 100%
+        Split chooses a >110% filled shard and reduces it to 100%
+        If the shard is between 90% and 110% full, nobody touches it
+        """
+        await self.build_shard_resources_index()
+        while True:
+            await self.wait_for_indexing()
+
+            shards = await self.get_rebalance_shards()
+
+            # Any shards to split?
+            shard_to_split = next((s for s in shards[::-1] if needs_split(s)), None)
+            if shard_to_split is not None:
+                await self.split_shard(shard_to_split, shards)
+                continue
+
+            # Any shards to merge?
+            shard_to_merge = next((s for s in shards if needs_merge(s, shards)), None)
+            if shard_to_merge is not None:
+                await self.merge_shard(shard_to_merge, shards)
+            else:
+                break
+
+    async def split_shard(self, shard_to_split: RebalanceShard, shards: list[RebalanceShard]):
+        logger.info(
+            "Splitting excess of paragraphs to other shards",
+            extra={
+                "kbid": self.kbid,
+                "shard": shard_to_split.to_dict(),
+            },
+        )
 
-        shard_paragraphs = await get_shards_paragraphs(kbid)
-        total_paragraphs = sum([c for _, c in shard_paragraphs])
-
-        if (total_paragraphs / len(kb_shards.shards)) > (
-            settings.max_shard_paragraphs * 0.9  # 90% of the max
-        ):
-            # create new shard
-            async with datamanagers.with_rw_transaction() as txn:
-                kb_config = await datamanagers.kb.get_config(txn, kbid=kbid)
+        # First off, calculate if the excess fits in the other shards or we need to add a new shard.
+        # Note that we don't filter out the active shard on purpose.
+        excess = shard_to_split.paragraphs - settings.max_shard_paragraphs
+        other_shards = [s for s in shards if s.id != shard_to_split.id]
+        other_shards_capacity = sum(
+            [max(0, (settings.max_shard_paragraphs - s.paragraphs)) for s in other_shards]
+        )
+        if excess > other_shards_capacity:
+            shards_to_add = math.ceil((excess - other_shards_capacity) / settings.max_shard_paragraphs)
+            logger.info(
+                "More shards needed",
+                extra={
+                    "kbid": self.kbid,
+                    "shards_to_add": shards_to_add,
+                    "all_shards": [s.to_dict() for s in shards],
+                },
+            )
+            # Add new shards where to rebalance the excess of paragraphs
+            async with (
+                locking.distributed_lock(locking.NEW_SHARD_LOCK.format(kbid=self.kbid)),
+                datamanagers.with_rw_transaction() as txn,
+            ):
+                kb_config = await datamanagers.kb.get_config(txn, kbid=self.kbid)
                 prewarm = kb_config is not None and kb_config.prewarm_enabled
-
                 sm = get_shard_manager()
-                await sm.create_shard_by_kbid(txn, kbid, prewarm_enabled=prewarm)
+                for _ in range(shards_to_add):
+                    await sm.create_shard_by_kbid(txn, self.kbid, prewarm_enabled=prewarm)
                 await txn.commit()
 
+            # Recalculate after having created shards, the active shard is a different one
+            shards = await self.get_rebalance_shards()
+
+        # Now, move resources to other shards as long as we are still over the max
+        for _ in range(MAX_MOVES_PER_SHARD):
+            shard_paragraphs = next(s.paragraphs for s in shards if s.id == shard_to_split.id)
+            excess = shard_paragraphs - settings.max_shard_paragraphs
+            if excess <= 0:
+                logger.info(
+                    "Shard rebalanced successfuly",
+                    extra={"kbid": self.kbid, "shard": shard_to_split.to_dict()},
+                )
+                break
+
+            target_shard, target_capacity = get_target_shard(shards, shard_to_split, skip_active=False)
+            if target_shard is None:
+                logger.warning("No target shard found for splitting", extra={"kbid": self.kbid})
+                break
 
-async def move_set_of_kb_resources(
-    context: ApplicationContext,
-    kbid: str,
-    from_shard_id: str,
-    to_shard_id: str,
-    count: int = 20,
-) -> None:
-    async with datamanagers.with_ro_transaction() as txn:
-        kb_shards = await datamanagers.cluster.get_kb_shards(txn, kbid=kbid)
-    if kb_shards is None:  # pragma: no cover
-        logger.warning("No shards found for kb. This should not happen.", extra={"kbid": kbid})
-        return
+            moved_paragraphs = await self.move_paragraphs(
+                from_shard=shard_to_split,
+                to_shard=target_shard,
+                max_paragraphs=min(excess, target_capacity),
+            )
 
-    logger.info(
-        "Rebalancing kb shards",
-        extra={"kbid": kbid, "from": from_shard_id, "to": to_shard_id, "count": count},
-    )
+            # Update shard paragraph counts
+            shard_to_split.paragraphs -= moved_paragraphs
+            target_shard.paragraphs += moved_paragraphs
+            shards.sort(key=lambda x: x.paragraphs)
 
-    from_shard = [s for s in kb_shards.shards if s.shard == from_shard_id][0]
-    to_shard = [s for s in kb_shards.shards if s.shard == to_shard_id][0]
+            await self.wait_for_indexing()
 
-    request = nodereader_pb2.SearchRequest(
-        shard=from_shard.nidx_shard_id,
-        paragraph=False,
-        document=True,
-        result_per_page=count,
-    )
-    request.field_filter.field.field_type = "a"
-    request.field_filter.field.field_id = "title"
-    search_response: nodereader_pb2.SearchResponse = await get_nidx_searcher_client().Search(request)
+    async def merge_shard(self, shard_to_merge: RebalanceShard, shards: list[RebalanceShard]):
+        logger.info(
+            "Merging shard",
+            extra={
+                "kbid": self.kbid,
+                "shard": shard_to_merge.to_dict(),
+            },
+        )
+        empty_shard = False
+
+        for _ in range(MAX_MOVES_PER_SHARD):
+            resources_count = len(self.index.get(shard_to_merge.id, []))
+            if resources_count == 0:
+                logger.info(
+                    "Shard is now empty",
+                    extra={
+                        "kbid": self.kbid,
+                        "shard": shard_to_merge.to_dict(),
+                    },
+                )
+                empty_shard = True
+                break
+
+            logger.info(
+                "Shard not yet empty",
+                extra={
+                    "kbid": self.kbid,
+                    "shard": shard_to_merge.to_dict(),
+                    "remaining": resources_count,
+                },
+            )
 
-    for result in search_response.document.results:
-        resource_id = result.uuid
-        try:
-            async with (
-                datamanagers.with_transaction() as txn,
-                locking.distributed_lock(
-                    locking.RESOURCE_INDEX_LOCK.format(kbid=kbid, resource_id=resource_id)
-                ),
-            ):
-                found_shard_id = await datamanagers.resources.get_resource_shard_id(
-                    txn, kbid=kbid, rid=resource_id, for_update=True
+            target_shard, target_capacity = get_target_shard(shards, shard_to_merge, skip_active=True)
+            if target_shard is None:
+                logger.warning(
+                    "No target shard could be found for merging. Moving on",
+                    extra={"kbid": self.kbid, "shard": shard_to_merge.to_dict()},
                 )
-                if found_shard_id is None:
-                    # resource deleted
-                    continue
-                if found_shard_id != from_shard_id:
-                    # resource could have already been moved
-                    continue
+                break
 
-                await datamanagers.resources.set_resource_shard_id(
-                    txn, kbid=kbid, rid=resource_id, shard=to_shard_id
+            moved_paragraphs = await self.move_paragraphs(
+                from_shard=shard_to_merge,
+                to_shard=target_shard,
+                max_paragraphs=target_capacity,
+            )
+
+            # Update shard paragraph counts
+            shard_to_merge.paragraphs -= moved_paragraphs
+            target_shard.paragraphs += moved_paragraphs
+            shards.sort(key=lambda x: x.paragraphs)
+
+            await self.wait_for_indexing()
+
+        if empty_shard:
+            # Build the index again, and make sure there is no resource assigned to this shard
+            await self.build_shard_resources_index()
+            shard_resources = self.index.get(shard_to_merge.id, set())
+            if len(shard_resources) > 0:
+                logger.error(
+                    f"Shard expected to be empty, but it isn't. Won't be deleted.",
+                    extra={
+                        "kbid": self.kbid,
+                        "shard": shard_to_merge.id,
+                        "resources": list(shard_resources)[:30],
+                    },
                 )
-                await index_resource_to_shard(context, kbid, resource_id, to_shard)
-                await delete_resource_from_shard(context, kbid, resource_id, from_shard)
-                await txn.commit()
+                return
+
+            # If shard was emptied, delete it
+            async with locking.distributed_lock(locking.NEW_SHARD_LOCK.format(kbid=self.kbid)):
+                async with datamanagers.with_rw_transaction() as txn:
+                    kb_shards = await datamanagers.cluster.get_kb_shards(
+                        txn, kbid=self.kbid, for_update=True
+                    )
+                    if kb_shards is not None:
+                        logger.info(
+                            "Deleting empty shard",
+                            extra={
+                                "kbid": self.kbid,
+                                "shard_id": shard_to_merge.id,
+                                "nidx_shard_id": shard_to_merge.nidx_id,
+                            },
+                        )
+
+                        # Delete shards from kb shards in maindb
+                        to_delete, to_delete_idx = next(
+                            (s, idx)
+                            for idx, s in enumerate(kb_shards.shards)
+                            if s.shard == shard_to_merge.id
+                        )
+                        kb_shards.shards.remove(to_delete)
+                        if to_delete_idx <= kb_shards.actual:
+                            # Only decrement the actual pointer if we remove before the pointer.
+                            kb_shards.actual -= 1
+                        assert kb_shards.actual >= 0
+                        await datamanagers.cluster.update_kb_shards(
+                            txn, kbid=self.kbid, shards=kb_shards
+                        )
+                        await txn.commit()
+
+                # Delete shard from nidx
+                await get_nidx_api_client().DeleteShard(
+                    noderesources_pb2.ShardId(id=to_delete.nidx_shard_id)
+                )
+
+
+async def get_resource_paragraphs_count(resource_id: str, nidx_shard_id: str) -> int:
+    # Do a search on the fields (paragraph) index and return the number of paragraphs this resource has
+    try:
+        request = nodereader_pb2.SearchRequest(
+            shard=nidx_shard_id,
+            paragraph=True,
+            document=False,
+            result_per_page=0,
+            field_filter=nodereader_pb2.FilterExpression(
+                resource=nodereader_pb2.FilterExpression.ResourceFilter(resource_id=resource_id)
+            ),
+        )
+        search_response: nodereader_pb2.SearchResponse = await get_nidx_searcher_client().Search(request)
+        return search_response.paragraph.total
+    except AioRpcError as exc:  # pragma: no cover
+        if exc.code() == StatusCode.NOT_FOUND:
+            logger.warning(f"Shard not found in nidx", extra={"nidx_shard_id": nidx_shard_id})
+            return 0
+        raise
+
+
+def get_target_shard(
+    shards: list[RebalanceShard], rebalanced_shard: RebalanceShard, skip_active: bool = True
+) -> tuple[Optional[RebalanceShard], int]:
+    """
+    Return the biggest shard with capacity (< 90% of the max paragraphs per shard).
+    """
+    target_shard = next(
+        reversed(
+            [
+                s
+                for s in shards
+                if s.id != rebalanced_shard.id
+                and s.paragraphs < settings.max_shard_paragraphs * 0.9
+                and (not skip_active or (skip_active and not s.active))
+            ]
+        ),
+        None,
+    )
+    if target_shard is None:  # pragma: no cover
+        return None, 0
+
+    # Aim to fill target shards up to 100% of max
+    capacity = int(max(0, settings.max_shard_paragraphs - target_shard.paragraphs))
+    return target_shard, capacity
+
+
+async def get_shard_paragraph_count(nidx_shard_id: str) -> int:
+    # Do a search on the fields (paragraph) index
+    try:
+        request = nodereader_pb2.SearchRequest(
+            shard=nidx_shard_id,
+            paragraph=True,
+            document=False,
+            result_per_page=0,
+        )
+        search_response: nodereader_pb2.SearchResponse = await get_nidx_searcher_client().Search(request)
+        return search_response.paragraph.total
+    except AioRpcError as exc:  # pragma: no cover
+        if exc.code() == StatusCode.NOT_FOUND:
+            logger.warning(f"Shard not found in nidx", extra={"nidx_shard_id": nidx_shard_id})
+            return 0
+        raise
+
+
+async def move_resource_to_shard(
+    context: ApplicationContext,
+    kbid: str,
+    resource_id: str,
+    from_shard: writer_pb2.ShardObject,
+    to_shard: writer_pb2.ShardObject,
+) -> bool:
+    indexed_to_new = False
+    deleted_from_old = False
+    try:
+        async with (
+            datamanagers.with_transaction() as txn,
+            locking.distributed_lock(
+                locking.RESOURCE_INDEX_LOCK.format(kbid=kbid, resource_id=resource_id)
+            ),
+        ):
+            found_shard_id = await datamanagers.resources.get_resource_shard_id(
+                txn, kbid=kbid, rid=resource_id, for_update=True
+            )
+            if found_shard_id is None:  # pragma: no cover
+                # resource deleted
+                return False
+            if found_shard_id != from_shard.shard:  # pragma: no cover
+                # resource could have already been moved
+                return False
+
+            await datamanagers.resources.set_resource_shard_id(
+                txn, kbid=kbid, rid=resource_id, shard=to_shard.shard
+            )
+            await index_resource_to_shard(context, kbid, resource_id, to_shard)
+            indexed_to_new = True
+            await delete_resource_from_shard(context, kbid, resource_id, from_shard)
+            deleted_from_old = True
+            await txn.commit()
+            return True
+    except Exception:
+        logger.exception(
+            "Failed to move resource",
+            extra={"kbid": kbid, "resource_id": resource_id},
+        )
+        # XXX Not ideal failure situation here. Try reverting the whole move even though it could be redundant
+        try:
+            if indexed_to_new:
+                await delete_resource_from_shard(context, kbid, resource_id, to_shard)
+            if deleted_from_old:
+                await index_resource_to_shard(context, kbid, resource_id, from_shard)
         except Exception:
             logger.exception(
-                "Failed to move resource",
+                "Failed to revert move resource. Hopefully you never see this message.",
                 extra={"kbid": kbid, "resource_id": resource_id},
             )
-            # XXX Not ideal failure situation here. Try reverting the whole move even though it could be redundant
-            try:
-                await index_resource_to_shard(context, kbid, resource_id, from_shard)
-                await delete_resource_from_shard(context, kbid, resource_id, to_shard)
-            except Exception:
-                logger.exception(
-                    "Failed to revert move resource. Hopefully you never see this message.",
-                    extra={"kbid": kbid, "resource_id": resource_id},
-                )
+        return False
 
 
-async def rebalance_kb(context: ApplicationContext, kbid: str) -> None:
-    await maybe_add_shard(kbid)
+def needs_split(shard: RebalanceShard) -> bool:
+    """
+    Return true if the shard is more than 110% of the max.
+
+    Active shards are not considered for splitting: the shard creator subscriber will
+    eventually create a new shard, make it the active one and the previous one, if
+    too full, will be split.
+    """
+    return not shard.active and (shard.paragraphs > (settings.max_shard_paragraphs * 1.1))
 
-    shard_paragraphs = await get_shards_paragraphs(kbid)
-    rebalanced_shards = set()
-    while any(paragraphs > settings.max_shard_paragraphs for _, paragraphs in shard_paragraphs):
-        # find the shard with the least/most paragraphs
-        smallest_shard = shard_paragraphs[0][0]
-        largest_shard = shard_paragraphs[-1][0]
-        assert smallest_shard != largest_shard
 
-        if smallest_shard in rebalanced_shards:
-            # XXX This is to prevent flapping data between shards on a single pass
-            # if we already rebalanced this shard, then we can't do anything else
-            break
+def needs_merge(shard: RebalanceShard, all_shards: list[RebalanceShard]) -> bool:
+    """
+    Returns true if a shard is less 75% full and there is enough capacity on the other shards to fit it.
 
-        await move_set_of_kb_resources(context, kbid, largest_shard, smallest_shard)
+    Active shards are not considered for merging. Shards that are more than 75% full are also skipped.
+    """
+    if shard.active:
+        return False
+    if shard.paragraphs > (settings.max_shard_paragraphs * 0.75):
+        return False
+    other_shards = [s for s in all_shards if s.id != shard.id and not s.active]
+    other_shards_capacity = sum(
+        [max(0, (settings.max_shard_paragraphs - s.paragraphs)) for s in other_shards]
+    )
+    return shard.paragraphs < other_shards_capacity
 
-        rebalanced_shards.add(largest_shard)
 
-        shard_paragraphs = await get_shards_paragraphs(kbid)
+async def rebalance_kb(context: ApplicationContext, kbid: str) -> None:
+    rebalancer = Rebalancer(context, kbid)
+    await rebalancer.rebalance_shards()
 
 
 async def run(context: ApplicationContext) -> None:
@@ -185,8 +514,12 @@ async def run(context: ApplicationContext) -> None:
             # get all kb ids
             async with datamanagers.with_ro_transaction() as txn:
                 kbids = [kbid async for kbid, _ in datamanagers.kb.get_kbs(txn)]
-            # go through each kb and see if shards need to be reduced in size
+            # go through each kb and see if shards need to be rebalanced
             for kbid in kbids:
+                if not has_feature(
+                    const.Features.REBALANCE_ENABLED, default=False, context={"kbid": kbid}
+                ):
+                    continue
                 async with locking.distributed_lock(locking.KB_SHARDS_LOCK.format(kbid=kbid)):
                     await rebalance_kb(context, kbid)
     except locking.ResourceLocked as exc:

nucliadb/common/cluster/rollover.py

@@ -47,6 +47,7 @@ from .utils import (
     get_resource,
     get_rollover_resource_index_message,
     index_resource_to_shard,
+    wait_for_nidx,
 )
 
 logger = logging.getLogger(__name__)
@@ -256,6 +257,7 @@ async def index_to_rollover_index(
             for rid in resource_ids
         ]
         await asyncio.gather(*batch)
+        await wait_for_indexing_to_catch_up(app_context)
 
     async with datamanagers.with_transaction() as txn:
         state.resources_indexed = True
@@ -264,6 +266,22 @@ async def index_to_rollover_index(
         await txn.commit()
 
 
+async def wait_for_indexing_to_catch_up(app_context: ApplicationContext):
+    try:
+        app_context.nats_manager
+    except AssertionError:
+        logger.warning("Nats manager not initialized. Cannot wait for indexing to catch up")
+        return
+    max_pending = 1000
+    while True:
+        try:
+            await wait_for_nidx(app_context.nats_manager, max_wait_seconds=60, max_pending=max_pending)
+            return
+        except asyncio.TimeoutError:
+            logger.warning(f"Nidx is behind more than {max_pending} messages. Throttling rollover.")
+            await asyncio.sleep(30)
+
+
 async def _index_resource_to_rollover_index(
     app_context: ApplicationContext,
     rollover_shards: writer_pb2.Shards,

nucliadb/common/cluster/utils.py

@@ -32,6 +32,7 @@ from nucliadb.common.cluster.settings import settings
 from nucliadb.ingest.orm import index_message
 from nucliadb.ingest.orm.resource import Resource
 from nucliadb_protos import writer_pb2
+from nucliadb_utils.nats import NatsConnectionManager
 from nucliadb_utils.utilities import Utility, clean_utility, get_utility, set_utility
 
 if TYPE_CHECKING:  # pragma: no cover
@@ -125,3 +126,28 @@ async def delete_resource_from_shard(
     partition = partitioning.generate_partition(kbid, resource_id)
 
     await sm.delete_resource(shard, resource_id, 0, str(partition), kbid)
+
+
+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
+
+
+async def wait_for_nidx(
+    nats_manager: NatsConnectionManager,
+    max_pending: int,
+    poll_interval_seconds: int = 5,
+    max_wait_seconds: int = 60,
+):
+    async with asyncio.timeout(max_wait_seconds):  # type: ignore
+        while True:
+            pending = await get_nats_consumer_pending_messages(
+                nats_manager, stream="nidx", consumer="nidx"
+            )
+            if pending < max_pending:
+                return
+            await asyncio.sleep(poll_interval_seconds)

nucliadb/ingest/orm/processor/processor.py

@@ -446,26 +446,27 @@ class Processor:
             # a resource was move to another shard while it was being indexed
             shard_id = await datamanagers.resources.get_resource_shard_id(txn, kbid=kbid, rid=uuid)
 
-        shard = None
-        if shard_id is not None:
-            # Resource already has a shard assigned
-            shard = await kb.get_resource_shard(shard_id)
-            if shard is None:
-                raise AttributeError("Shard not available")
-        else:
-            # It's a new resource, get KB's current active shard to place new resource on
-            shard = await self.index_node_shard_manager.get_current_active_shard(txn, kbid)
-            if shard is None:
-                # No current shard available, create a new one
-                kb_config = await datamanagers.kb.get_config(txn, kbid=kbid)
-                prewarm = kb_config is not None and kb_config.prewarm_enabled
-                shard = await self.index_node_shard_manager.create_shard_by_kbid(
-                    txn, kbid, prewarm_enabled=prewarm
+            shard = None
+            if shard_id is not None:
+                # Resource already has a shard assigned
+                shard = await kb.get_resource_shard(shard_id)
+                if shard is None:
+                    raise AttributeError("Shard not available")
+            else:
+                # It's a new resource, get KB's current active shard to place new resource on
+                shard = await self.index_node_shard_manager.get_current_active_shard(txn, kbid)
+                if shard is None:
+                    # No current shard available, create a new one
+                    async with locking.distributed_lock(locking.NEW_SHARD_LOCK.format(kbid=kbid)):
+                        kb_config = await datamanagers.kb.get_config(txn, kbid=kbid)
+                        prewarm = kb_config is not None and kb_config.prewarm_enabled
+                        shard = await self.index_node_shard_manager.create_shard_by_kbid(
+                            txn, kbid, prewarm_enabled=prewarm
+                        )
+                await datamanagers.resources.set_resource_shard_id(
+                    txn, kbid=kbid, rid=uuid, shard=shard.shard
                 )
-            await datamanagers.resources.set_resource_shard_id(
-                txn, kbid=kbid, rid=uuid, shard=shard.shard
-            )
-        return shard
+            return shard
 
     @processor_observer.wrap({"type": "index_resource"})
     async def index_resource(

nucliadb/purge/__init__.py

@@ -19,6 +19,7 @@
 #
 import asyncio
 import importlib.metadata
+from itertools import batched  # type: ignore
 from typing import AsyncGenerator
 
 from nucliadb.common import datamanagers
@@ -233,7 +234,7 @@ async def purge_kb_vectorsets(driver: Driver, storage: Storage):
                     fields.extend((await resource.get_fields(force=True)).values())
 
             logger.info(f"Purging {len(fields)} fields for vectorset {vectorset}", extra={"kbid": kbid})
-            for fields_batch in batchify(fields, 20):
+            for fields_batch in batched(fields, n=20):
                 tasks = []
                 for field in fields_batch:
                     if purge_payload.storage_key_kind == VectorSetConfig.StorageKeyKind.UNSET:
@@ -317,9 +318,3 @@ def run() -> int:  # pragma: no cover
     setup_logging()
     errors.setup_error_handling(importlib.metadata.distribution("nucliadb").version)
     return asyncio.run(main())
-
-
-def batchify(iterable, n=1):
-    """Yield successive n-sized chunks from iterable."""
-    for i in range(0, len(iterable), n):
-        yield iterable[i : i + n]

{nucliadb-6.9.1.post5192.dist-info → nucliadb-6.9.1.post5208.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nucliadb
-Version: 6.9.1.post5192
+Version: 6.9.1.post5208
 Summary: NucliaDB
 Author-email: Nuclia <nucliadb@nuclia.com>
 License-Expression: AGPL-3.0-or-later
@@ -19,11 +19,11 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3 :: Only
 Requires-Python: <4,>=3.9
 Description-Content-Type: text/markdown
-Requires-Dist: nucliadb-telemetry[all]>=6.9.1.post5192
-Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.9.1.post5192
-Requires-Dist: nucliadb-protos>=6.9.1.post5192
-Requires-Dist: nucliadb-models>=6.9.1.post5192
-Requires-Dist: nidx-protos>=6.9.1.post5192
+Requires-Dist: nucliadb-telemetry[all]>=6.9.1.post5208
+Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.9.1.post5208
+Requires-Dist: nucliadb-protos>=6.9.1.post5208
+Requires-Dist: nucliadb-models>=6.9.1.post5208
+Requires-Dist: nidx-protos>=6.9.1.post5208
 Requires-Dist: nucliadb-admin-assets>=1.0.0.post1224
 Requires-Dist: nuclia-models>=0.50.0
 Requires-Dist: uvicorn[standard]

{nucliadb-6.9.1.post5192.dist-info → nucliadb-6.9.1.post5208.dist-info}/RECORD

@@ -86,10 +86,10 @@ nucliadb/common/cluster/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIX
 nucliadb/common/cluster/exceptions.py,sha256=t7v_l93t44l2tQpdQXgO_w-c4YZRcaayOz1A2i0w4RQ,1258
 nucliadb/common/cluster/grpc_node_dummy.py,sha256=JkufazWzMA4KFEU8EBkMbiiDW4C8lLcRhiiCxP7aCQY,2949
 nucliadb/common/cluster/manager.py,sha256=p-haaGEnCa-20t-I2XEo4QJ5ZC1QgJ6p2jzXFYVB6nQ,12346
-nucliadb/common/cluster/rebalance.py,sha256=U-LgmDZq7JaEKh4fruUoW8VrqAAnpQrdwqN_p0tbCKo,8899
-nucliadb/common/cluster/rollover.py,sha256=i_W5_ds_5c_SZWxQGg4I8c_d7tP2VEPSnCVq5HU55jI,26490
+nucliadb/common/cluster/rebalance.py,sha256=s0YJ07Y358T9x22QnHHEvhmHqf1CnQWQD6oJznOo2Xc,22239
+nucliadb/common/cluster/rollover.py,sha256=kmVCdyjJ1dilnSodHMqf0EUkTjPC5H0aA4JqW9KsEa4,27168
 nucliadb/common/cluster/settings.py,sha256=f6Y5K0PGahkedwe5wtkWMnbqwDFJgOOwX_MOIGwH9Dg,2271
-nucliadb/common/cluster/utils.py,sha256=IfW5PA7Ab26xWUYNOc3yBoksWV1GK4BGhTRo1vnHNKg,4662
+nucliadb/common/cluster/utils.py,sha256=E4GqidwTKczJX_lTnncBCof2fL4CFVVF1eLiz9NWjlc,5494
 nucliadb/common/cluster/standalone/__init__.py,sha256=itSI7dtTwFP55YMX4iK7JzdMHS5CQVUiB1XzQu4UBh8,833
 nucliadb/common/cluster/standalone/utils.py,sha256=af3r-x_GF7A6dwIAhZLR-r-SZQEVxsFrDKeMfUTA6G0,1908
 nucliadb/common/context/__init__.py,sha256=IKAHuiCjbOEsqfLozWwJ6mRFzFncsZMyxNC5E_XZ5EM,6016
@@ -174,7 +174,7 @@ nucliadb/ingest/orm/utils.py,sha256=fCQRuyecgqhaY7mcBG93oaXMkzkKb9BFjOcy4-ZiSNw,
 nucliadb/ingest/orm/processor/__init__.py,sha256=xhDNKCxY0XNOlIVKEtM8QT75vDUkJIt7K-_VgGbbOQU,904
 nucliadb/ingest/orm/processor/auditing.py,sha256=gxn5v30KVaH0TnIjo715mWjzKGJ-DMviElEXJG9BNN4,4612
 nucliadb/ingest/orm/processor/data_augmentation.py,sha256=v-pj4GbBWSuO8dQyahs5UDr5ghsyfhCZDS0ftKd6ZYc,5179
-nucliadb/ingest/orm/processor/processor.py,sha256=t9k0AxKRC3w5btv_e4Xw93vMwWCtEytWDV1kxFXCtbk,31609
+nucliadb/ingest/orm/processor/processor.py,sha256=S7myNucnK8v_5y4gZxS_5nL-glGx_fZe2l8gSzDy2CU,31808
 nucliadb/ingest/orm/processor/sequence_manager.py,sha256=kUH0bCuM6NqpA0xSwfyb9igig3Btu57pc8VYnKggqx4,1693
 nucliadb/ingest/service/__init__.py,sha256=LHQFUkdmNBOWqBG0Md9sMMI7g5TQZ-hLAnhw6ZblrJg,2002
 nucliadb/ingest/service/exceptions.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
@@ -193,7 +193,7 @@ nucliadb/models/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,8
 nucliadb/models/responses.py,sha256=qnuOoc7TrVSUnpikfTwHLKez47_DE4mSFzpxrwtqijA,1599
 nucliadb/models/internal/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
 nucliadb/models/internal/processing.py,sha256=F_WpDItHsBWuMNeLMxXOKqpZoRYdjPuB2l0s29PVcfI,4213
-nucliadb/purge/__init__.py,sha256=FWqY2ln2ecGgVcZaMhfs8nTt4HzpprubkCsurSrIrC8,13070
+nucliadb/purge/__init__.py,sha256=RFqMj3JRML2sGrQtmRm2fzRg350f0PhkBpZh-KRWnlc,12954
 nucliadb/purge/orphan_shards.py,sha256=5QvTBoYJ2h14TUvJmZxphHJPf_WB-VVC9453n-zjSrU,7747
 nucliadb/reader/__init__.py,sha256=C5Efic7WlGm2U2C5WOyquMFbIj2Pojwe_8mwzVYnOzE,1304
 nucliadb/reader/app.py,sha256=Se-BFTE6d1v1msLzQn4q5XIhjnSxa2ckDSHdvm7NRf8,3096
@@ -385,8 +385,8 @@ nucliadb/writer/tus/local.py,sha256=7jYa_w9b-N90jWgN2sQKkNcomqn6JMVBOVeDOVYJHto,
 nucliadb/writer/tus/s3.py,sha256=vu1BGg4VqJ_x2P1u2BxqPKlSfw5orT_a3R-Ln5oPUpU,8483
 nucliadb/writer/tus/storage.py,sha256=ToqwjoYnjI4oIcwzkhha_MPxi-k4Jk3Lt55zRwaC1SM,2903
 nucliadb/writer/tus/utils.py,sha256=MSdVbRsRSZVdkaum69_0wku7X3p5wlZf4nr6E0GMKbw,2556
-nucliadb-6.9.1.post5192.dist-info/METADATA,sha256=dS74zqOkgFw6qgTfAPpcIj1yfslFWNp1X9233WPKPMk,4158
-nucliadb-6.9.1.post5192.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-nucliadb-6.9.1.post5192.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
-nucliadb-6.9.1.post5192.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
-nucliadb-6.9.1.post5192.dist-info/RECORD,,
+nucliadb-6.9.1.post5208.dist-info/METADATA,sha256=SyOTfLr01715WH-X7Zg6ejg73H-4ImjqnB8hwg9QpPs,4158
+nucliadb-6.9.1.post5208.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+nucliadb-6.9.1.post5208.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
+nucliadb-6.9.1.post5208.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
+nucliadb-6.9.1.post5208.dist-info/RECORD,,