Catch nats timeout errors and handle them better (#2630)

* Catch nats timeout errors

* Add exception that we control at processing

* Add another one

nucliadb/src/nucliadb/ingest/orm/processor/__init__.py

@@ -22,6 +22,8 @@
 from typing import Optional
 
 import aiohttp.client_exceptions
+import nats.errors
+import nats.js.errors
 
 from nucliadb.common import datamanagers, locking
 from nucliadb.common.cluster.settings import settings as cluster_settings
@@ -325,6 +327,8 @@ async def txn(
                 aiohttp.client_exceptions.ClientError,
                 ConflictError,
                 MaindbServerError,
+                nats.errors.NoRespondersError,
+                nats.js.errors.NoStreamResponseError,
             ):  # pragma: no cover
                 # Unhandled exceptions here that should bubble and hard fail
                 # XXX We swallow too many exceptions here!

nucliadb/src/nucliadb/writer/api/v1/transaction.py

@@ -21,7 +21,11 @@
 from fastapi import HTTPException
 
 from nucliadb_protos.writer_pb2 import BrokerMessage
-from nucliadb_utils.transaction import MaxTransactionSizeExceededError, TransactionCommitTimeoutError
+from nucliadb_utils.transaction import (
+    MaxTransactionSizeExceededError,
+    StreamingServerTimeoutError,
+    TransactionCommitTimeoutError,
+)
 from nucliadb_utils.utilities import get_transaction_utility
 
 
@@ -39,3 +43,8 @@ async def commit(writer: BrokerMessage, partition: int, wait: bool = True) -> No
             status_code=413,
             detail="Transaction size exceeded. The resource is too large to be stored. Consider using file fields or split into multiple requests.",
         )
+    except StreamingServerTimeoutError:
+        raise HTTPException(
+            status_code=504,
+            detail="Timeout waiting for the streaming server to respond. Please back off and retry.",
+        )

nucliadb_utils/src/nucliadb_utils/transaction.py

@@ -58,6 +58,10 @@ class MaxTransactionSizeExceededError(Exception):
     pass
 
 
+class StreamingServerTimeoutError(Exception):
+    pass
+
+
 class LocalTransactionUtility:
     async def commit(
         self,
@@ -204,6 +208,16 @@ async def commit(
             res = await self.js.publish(target_subject, writer.SerializeToString(), headers=headers)
         except nats.errors.MaxPayloadError as ex:
             raise MaxTransactionSizeExceededError() from ex
+        except nats.errors.TimeoutError as ex:
+            logger.exception(
+                "Nats server timeout error on publish",
+                extra={
+                    "partition": partition,
+                    "kbid": writer.kbid,
+                    "resource": writer.uuid,
+                },
+            )
+            raise StreamingServerTimeoutError() from ex
 
         waiting_for.seq = res.seq