nucliadb-utils 6.1.0.post2589__py3-none-any.whl → 6.1.0.post2602__py3-none-any.whl

nucliadb_utils/nats.py

@@ -19,6 +19,7 @@
 
 import asyncio
 import logging
+import sys
 import time
 from functools import cached_property
 from typing import Any, Awaitable, Callable, Optional, Union
@@ -126,6 +127,9 @@ class NatsConnectionManager:
         self._lock = asyncio.Lock()
         self._healthy = True
         self._last_unhealthy: Optional[float] = None
+        self._needs_reconnection = False
+        self._reconnect_task: Optional[asyncio.Task] = None
+        self._expected_subscriptions: set[str] = set()
 
     def healthy(self) -> bool:
         if not self._healthy:
@@ -159,8 +163,23 @@ class NatsConnectionManager:
         async with self._lock:
             self._nc = await nats.connect(**options)
 
+        self._expected_subscription_task = asyncio.create_task(self._verify_expected_subscriptions())
+
     async def finalize(self):
         async with self._lock:
+            if self._reconnect_task:
+                self._reconnect_task.cancel()
+                try:
+                    await self._reconnect_task
+                except asyncio.CancelledError:
+                    pass
+
+            self._expected_subscription_task.cancel()
+            try:
+                await self._expected_subscription_task
+            except asyncio.CancelledError:
+                pass
+
             # Finalize push subscriptions
             for sub, _ in self._subscriptions:
                 try:
@@ -199,42 +218,75 @@ class NatsConnectionManager:
         logger.warning(
             f"Reconnected to NATS {self._nc.connected_url.netloc}. Attempting to re-subscribe."
         )
-        async with self._lock:
-            existing_subs = self._subscriptions
-            self._subscriptions = []
-            for sub, recon_callback in existing_subs:
-                try:
-                    await sub.drain()
-                    await recon_callback()
-                except Exception:
-                    logger.exception(
-                        f"Error resubscribing to {sub.subject} on {self._nc.connected_url.netloc}"
-                    )
-                    # should force exit here to restart the service
-                    self._healthy = False
-                    raise
-
-            existing_pull_subs = self._pull_subscriptions
-            self._pull_subscriptions = []
-            for pull_sub, task, recon_callback, cancelled in existing_pull_subs:
-                cancelled.set()
-                task.cancel()
-                try:
-                    await task
-                except asyncio.CancelledError:
-                    pass
-                try:
-                    await pull_sub.unsubscribe()
-                    await recon_callback()
-                except Exception:
-                    logger.exception(
-                        f"Error resubscribing to pull subscription on {self._nc.connected_url.netloc}"
-                    )
-                    # should force exit here to restart the service
-                    self._healthy = False
-                    raise
-        self._healthy = True
-        self._last_unhealthy = None  # reset the last unhealthy time
+        # This callback may be interrupted by NATS. In order to avoid being interrupted, we spawn our own task to do
+        # the reconnection. If we are interrupted we could lose subscriptions
+        self._needs_reconnection = True
+        if self._reconnect_task is None:
+            self._reconnect_task = asyncio.create_task(self._reconnect())
+
+    async def _reconnect(self):
+        tries = 0
+        # Loop in case we receive another reconnection request while one is ongoing
+        while self._needs_reconnection:
+            if tries > 5:
+                logger.error(
+                    "Too many consecutive reconnection to NATS. Something might be wrong. Exiting."
+                )
+                sys.exit(1)
+
+            try:
+                self._needs_reconnection = False
+
+                async with self._lock:
+                    existing_subs = self._subscriptions
+                    self._subscriptions = []
+                    for sub, recon_callback in existing_subs:
+                        try:
+                            await sub.drain()
+                            await recon_callback()
+                        except Exception:
+                            logger.exception(
+                                f"Error resubscribing to {sub.subject} on {self._nc.connected_url.netloc}"
+                            )
+                            # should force exit here to restart the service
+                            self._healthy = False
+                            raise
+
+                    existing_pull_subs = self._pull_subscriptions
+                    self._pull_subscriptions = []
+                    for pull_sub, task, recon_callback, cancelled in existing_pull_subs:
+                        cancelled.set()
+                        task.cancel()
+                        try:
+                            await task
+                        except asyncio.CancelledError:
+                            pass
+                        try:
+                            await pull_sub.unsubscribe()
+                            await recon_callback()
+                        except Exception:
+                            logger.exception(
+                                f"Error resubscribing to pull subscription on {self._nc.connected_url.netloc}"
+                            )
+                            # should force exit here to restart the service
+                            self._healthy = False
+                            raise
+                self._healthy = True
+                self._last_unhealthy = None  # reset the last unhealthy time
+            except Exception:
+                if self._needs_reconnection:
+                    logger.warning("Error reconnecting to NATS, will retry")
+                else:
+                    logger.exception("Error reconnecting to NATS. Exiting.")
+                    sys.exit(1)
+
+            if self._needs_reconnection:
+                logger.warning(
+                    "While reconnecting to NATS subscriptions, a reconnect request was received. Reconnecting again.",
+                )
+            tries += 1
+
+        self._reconnect_task = None
 
     async def error_cb(self, e):  # pragma: no cover
         logger.error(f"There was an error on consumer: {e}", exc_info=e)
@@ -314,6 +366,10 @@ class NatsConnectionManager:
 
         task = asyncio.create_task(consume(psub, subject), name=f"pull_subscribe_{subject}")
         self._pull_subscriptions.append((psub, task, subscription_lost_cb, cancelled))
+
+        if durable:
+            self._expected_subscriptions.add(durable)
+
         return psub
 
     async def _remove_subscription(
@@ -340,3 +396,23 @@ class NatsConnectionManager:
     async def unsubscribe(self, subscription: Union[Subscription, JetStreamContext.PullSubscription]):
         await subscription.unsubscribe()
         await self._remove_subscription(subscription)
+
+        if isinstance(subscription, JetStreamContext.PullSubscription):
+            self._expected_subscriptions.remove(subscription._consumer)  # type: ignore
+
+    async def _verify_expected_subscriptions(self):
+        failures = 0
+        while True:
+            await asyncio.sleep(30)
+
+            existing_subs = set(sub._consumer for sub, _, _, _ in self._pull_subscriptions)
+            missing_subs = self._expected_subscriptions - existing_subs
+            if missing_subs:
+                logger.warning(f"Some NATS subscriptions are missing {missing_subs}")
+                failures += 1
+            else:
+                failures = 0
+
+            if failures >= 3:
+                logger.warning("Some NATS subscriptions are missing for too long, exiting")
+                sys.exit(1)

nucliadb_utils/utilities.py

@@ -92,6 +92,8 @@ def get_utility(ident: Union[Utility, str]):
 
 
 def set_utility(ident: Union[Utility, str], util: Any):
+    if ident in MAIN:
+        logger.warning(f"Overwriting previously set utility {ident}: {MAIN[ident]} with {util}")
     MAIN[ident] = util
 
 
@@ -237,6 +239,10 @@ def get_ingest() -> WriterStub:
 
 
 def start_partitioning_utility() -> PartitionUtility:
+    util = get_utility(Utility.PARTITION)
+    if util is not None:
+        return util
+
     util = PartitionUtility(
         partitions=nuclia_settings.nuclia_partitions,
         seed=nuclia_settings.nuclia_hash_seed,

{nucliadb_utils-6.1.0.post2589.dist-info → nucliadb_utils-6.1.0.post2602.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nucliadb_utils
-Version: 6.1.0.post2589
+Version: 6.1.0.post2602
 Home-page: https://nuclia.com
 License: BSD
 Classifier: Development Status :: 4 - Beta
@@ -24,8 +24,8 @@ Requires-Dist: PyNaCl
 Requires-Dist: pyjwt>=2.4.0
 Requires-Dist: memorylru>=1.1.2
 Requires-Dist: mrflagly>=0.2.9
-Requires-Dist: nucliadb-protos>=6.1.0.post2589
-Requires-Dist: nucliadb-telemetry>=6.1.0.post2589
+Requires-Dist: nucliadb-protos>=6.1.0.post2602
+Requires-Dist: nucliadb-telemetry>=6.1.0.post2602
 Provides-Extra: cache
 Requires-Dist: redis>=4.3.4; extra == "cache"
 Requires-Dist: orjson>=3.6.7; extra == "cache"

{nucliadb_utils-6.1.0.post2589.dist-info → nucliadb_utils-6.1.0.post2602.dist-info}/RECORD

@@ -8,7 +8,7 @@ nucliadb_utils/featureflagging.py,sha256=dKp8u4g52KTcupKFJkdGYk3Z2pGQdi_q7IaoY_G
 nucliadb_utils/grpc.py,sha256=apu0uePnkGHCAT7GRQ9YZfRYyFj26kJ440i8jitbM3U,3314
 nucliadb_utils/helpers.py,sha256=nPw8yod3hP-pxq80VF8QC36s7ygSg0dBUdfI-LatvCs,1600
 nucliadb_utils/indexing.py,sha256=Luaqcar3CySpdYOFp6Q9Fyr8ZYwhYhaKRHQ_VGL78f8,2318
-nucliadb_utils/nats.py,sha256=SGBWtP0h3_9Ne6VEfZ2OLnWlEvzfAj63U252FrTITjY,11870
+nucliadb_utils/nats.py,sha256=YsSuhGKpZjyZXjVT7uAF6KjRMWvp1k7-SF9ddFXdVC0,15059
 nucliadb_utils/partition.py,sha256=jBgy4Hu5Iwn4gjbPPcthSykwf-qNx-GcLAIwbzPd1d0,1157
 nucliadb_utils/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 nucliadb_utils/run.py,sha256=Es0_Bu5Yc-LWczvwL6gzWqSwC85RjDCk-0oFQAJi9g4,1827
@@ -16,7 +16,7 @@ nucliadb_utils/settings.py,sha256=BOuYTh08cJe7UgBgU9ALHCJm6eDezbdPsdAXM4HGEgo,81
 nucliadb_utils/signals.py,sha256=JRNv2y9zLtBjOANBf7krGfDGfOc9qcoXZ6N1nKWS2FE,2674
 nucliadb_utils/store.py,sha256=kQ35HemE0v4_Qg6xVqNIJi8vSFAYQtwI3rDtMsNy62Y,890
 nucliadb_utils/transaction.py,sha256=GjM754TU9940pfQY3vxAQAYlQdBpqfwrO-_gYkSjWrI,7979
-nucliadb_utils/utilities.py,sha256=vO6NUQxFpljTo7y6TemoOewMsXLgKTJmVPTtCWnjsDk,15921
+nucliadb_utils/utilities.py,sha256=_5pNsXBknULGrTvV4hgIjsOJ0sSbFq4ygseH89C1gdk,16128
 nucliadb_utils/aiopynecone/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
 nucliadb_utils/aiopynecone/client.py,sha256=MPyHnDXwhukJr7U3CJh7BpsekfSuOkyM4g5b9LLtzc8,22941
 nucliadb_utils/aiopynecone/exceptions.py,sha256=fUErx3ceKQK1MUbOnYcZhIzpNe8UVAptZE9JIRDLXDE,4000
@@ -59,8 +59,8 @@ nucliadb_utils/tests/indexing.py,sha256=YW2QhkhO9Q_8A4kKWJaWSvXvyQ_AiAwY1VylcfVQ
 nucliadb_utils/tests/local.py,sha256=fXIBasrvdaFJM-sw2wk1_oiFzBcm9O10iCyC-OiXwY8,1914
 nucliadb_utils/tests/nats.py,sha256=xqpww4jZjTKY9oPGlJdDJG67L3FIBQsa9qDHxILR8r8,7687
 nucliadb_utils/tests/s3.py,sha256=pl-RJFjA4MH6iXkqhsh5g8gDuEhrYu1nPZ-laxlrMlE,3704
-nucliadb_utils-6.1.0.post2589.dist-info/METADATA,sha256=qrXQ8HGQZ1FFhwMMAbyclo6eGgBvxlFXssBqqMAtZ4I,2055
-nucliadb_utils-6.1.0.post2589.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-nucliadb_utils-6.1.0.post2589.dist-info/top_level.txt,sha256=fE3vJtALTfgh7bcAWcNhcfXkNPp_eVVpbKK-2IYua3E,15
-nucliadb_utils-6.1.0.post2589.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-nucliadb_utils-6.1.0.post2589.dist-info/RECORD,,
+nucliadb_utils-6.1.0.post2602.dist-info/METADATA,sha256=7CB6_mYBRnxvaeq3fcpYyAvr62mVWpZCyhdHcT_CvLY,2055
+nucliadb_utils-6.1.0.post2602.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
+nucliadb_utils-6.1.0.post2602.dist-info/top_level.txt,sha256=fE3vJtALTfgh7bcAWcNhcfXkNPp_eVVpbKK-2IYua3E,15
+nucliadb_utils-6.1.0.post2602.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+nucliadb_utils-6.1.0.post2602.dist-info/RECORD,,