🐛 Fix endorser Nats connection going stale

app/services/onboarding/util/register_issuer_did.py

@@ -1,4 +1,5 @@
 import asyncio
+import os
 from logging import Logger
 
 from aries_cloudcontroller import (
@@ -18,6 +19,8 @@
 )
 from shared import ACAPY_ENDORSER_ALIAS
 
+MAX_ATTEMPTS = int(os.getenv("WAIT_ISSUER_DID_MAX_ATTEMPTS", "15"))
+
 
 async def create_connection_with_endorser(
     *,
@@ -241,7 +244,7 @@ async def wait_issuer_did_transaction_endorsed(
     issuer_controller: AcaPyClient,
     issuer_connection_id: str,
     logger: Logger,
-    max_attempts: int = 15,
+    max_attempts: int = MAX_ATTEMPTS,
     retry_delay: float = 1.0,
 ) -> None:
     attempt = 0

app/tests/e2e/test_did_exchange.py

@@ -1,3 +1,4 @@
+import asyncio
 from typing import Optional
 
 import pytest
@@ -99,8 +100,9 @@
                 filter_map={"their_did": alice_did},
             )
         finally:
+            await asyncio.sleep(1)  # Short sleep assists in avoiding 500 error
             # Delete connection records:
             await alice_member_client.delete(
                 f"{CONNECTIONS_BASE_PATH}/{alice_connection_id}"
             )
 
@@ -171,6 +173,7 @@
             filter_map={"connection_id": faber_connection_id},
         )
     finally:
+        await asyncio.sleep(1)  # Short sleep assists in avoiding 500 error
         # Delete connection records:
         await alice_member_client.delete(
             f"{CONNECTIONS_BASE_PATH}/{alice_connection_id}"

docker-compose.yaml

@@ -230,7 +230,7 @@ services:
     env_file:
       - environments/endorser/endorser.env
     healthcheck:
-      test: ["CMD-SHELL", "curl -f http://localhost:3009/health || exit 1"]
+      test: ["CMD-SHELL", "curl -f http://localhost:3009/health/ready || exit 1"]
       interval: 30s
       timeout: 10s
       retries: 3

endorser/main.py

@@ -1,3 +1,4 @@
+import asyncio
 import os
 from contextlib import asynccontextmanager
 
@@ -58,7 +59,7 @@ async def scalar_html():
     )
 
 
-@app.get("/health")
+@app.get("/health/live")
 @inject
 async def health_check(
     endorsement_processor: EndorsementProcessor = Depends(
@@ -71,3 +72,33 @@ async def health_check(
         raise HTTPException(
             status_code=503, detail="One or more background tasks are not running."
         )
+
+
+@app.get("/health/ready")
+@inject
+async def health_ready(
+    endorsement_processor: EndorsementProcessor = Depends(
+        Provide[Container.endorsement_processor]
+    ),
+):
+    try:
+        jetstream_status = await asyncio.wait_for(
+            endorsement_processor.check_jetstream(), timeout=5.0
+        )
+    except asyncio.TimeoutError:
+        raise HTTPException(
+            status_code=503,
+            detail={"status": "not ready", "error": "JetStream health check timed out"},
+        )
+    except Exception as e:  # pylint: disable=W0718
+        raise HTTPException(
+            status_code=500, detail={"status": "error", "error": str(e)}
+        )
+
+    if jetstream_status["is_working"]:
+        return {"status": "ready", "jetstream": jetstream_status}
+    else:
+        raise HTTPException(
+            status_code=503,
+            detail={"status": "not ready", "jetstream": "JetStream not ready"},
+        )

endorser/services/endorsement_processor.py

@@ -4,6 +4,7 @@
 from aries_cloudcontroller import AcaPyClient
 from nats.errors import BadSubscriptionError, Error, TimeoutError
 from nats.js.client import JetStreamContext
+from nats.js.errors import FetchTimeoutError
 
 from endorser.util.endorsement import accept_endorsement, should_accept_endorsement
 from shared.constants import (
@@ -85,7 +86,7 @@ async def _process_endorsement_requests(self) -> NoReturn:
         subscription = await self._subscribe()
         while True:
             try:
-                messages = await subscription.fetch(batch=1, timeout=60)
+                messages = await subscription.fetch(batch=1, timeout=60, heartbeat=1)
                 for message in messages:
                     message_subject = message.subject
                     message_data = message.data.decode()
@@ -103,9 +104,13 @@ async def _process_endorsement_requests(self) -> NoReturn:
                         )
                     finally:
                         await message.ack()
-            except TimeoutError:
-                logger.trace("Timeout fetching messages continuing...")
+            except FetchTimeoutError:
+                logger.trace("FetchTimeoutError continuing...")
                 await asyncio.sleep(0.1)
+            except TimeoutError as e:
+                logger.warning("Timeout error fetching messages re-subscribing: {}", e)
+                await subscription.unsubscribe()
+                subscription = await self._subscribe()
             except Exception:  # pylint: disable=W0718
                 logger.exception("Unexpected error in endorsement processing loop")
                 await asyncio.sleep(2)
@@ -192,3 +197,17 @@ async def _subscribe(self) -> JetStreamContext.PullSubscription:
         logger.debug("Subscribed to NATS subject")
 
         return subscription
+
+    async def check_jetstream(self):
+        try:
+            account_info = await self.jetstream.account_info()
+            is_working = account_info.streams > 0
+            logger.trace("JetStream check completed. Is working: {}", is_working)
+            return {
+                "is_working": is_working,
+                "streams_count": account_info.streams,
+                "consumers_count": account_info.consumers,
+            }
+        except Exception:  # pylint: disable=W0718
+            logger.exception("Caught exception while checking jetstream status")
+            return {"is_working": False}

endorser/tests/test_endorser_processor.py

@@ -6,6 +6,7 @@
 from nats.aio.client import Client as NATS
 from nats.errors import BadSubscriptionError, Error, TimeoutError
 from nats.js.client import JetStreamContext
+from nats.js.errors import FetchTimeoutError
 
 from endorser.services.endorsement_processor import EndorsementProcessor
 from shared.constants import (
@@ -128,7 +129,7 @@ async def test_process_endorsement_requests_timeout(
     mock_nats_client.pull_subscribe.return_value = mock_subscription
 
     # Simulate a timeout, then a CancelledError to stop the loop
-    mock_subscription.fetch.side_effect = [TimeoutError, asyncio.CancelledError]
+    mock_subscription.fetch.side_effect = [FetchTimeoutError, asyncio.CancelledError]
 
     # Test
     with patch("asyncio.sleep") as mock_sleep:

endorser/tests/test_main.py

@@ -1,9 +1,11 @@
+import asyncio
 from unittest.mock import AsyncMock, MagicMock, Mock, patch
 
 import pytest
 from fastapi import FastAPI, HTTPException
 
-from endorser.main import app, app_lifespan, health_check
+from endorser.main import app, app_lifespan, health_check, health_ready
+from endorser.services.endorsement_processor import EndorsementProcessor
 
 
 def test_create_app():
@@ -14,8 +16,9 @@ def test_create_app():
     # Get all routes in app
     routes = [route.path for route in app.routes]
 
-    expected_routes = "/health"
-    assert expected_routes in routes
+    expected_routes = ["/health/live", "/health/ready", "/docs"]
+    for route in expected_routes:
+        assert route in routes
 
 
 @pytest.mark.anyio
@@ -63,3 +66,89 @@ async def test_health_check_unhealthy():
         await health_check(endorsement_processor=endorsement_processor_mock)
     assert exc_info.value.status_code == 503
     assert exc_info.value.detail == "One or more background tasks are not running."
+
+
+@pytest.mark.anyio
+async def test_health_ready_success():
+    endorsement_processor_mock = AsyncMock(spec=EndorsementProcessor)
+
+    endorsement_processor_mock.check_jetstream.return_value = {
+        "is_working": True,
+        "streams_count": 1,
+        "consumers_count": 1,
+    }
+
+    response = await health_ready(endorsement_processor=endorsement_processor_mock)
+
+    assert response == {
+        "status": "ready",
+        "jetstream": {"is_working": True, "streams_count": 1, "consumers_count": 1},
+    }
+
+
+@pytest.mark.anyio
+async def test_health_ready_jetstream_not_working():
+    endorsement_processor_mock = AsyncMock(spec=EndorsementProcessor)
+
+    endorsement_processor_mock.check_jetstream.return_value = {
+        "is_working": False,
+        "error": "No streams available",
+    }
+
+    with pytest.raises(HTTPException) as exc_info:
+        await health_ready(endorsement_processor=endorsement_processor_mock)
+
+    assert exc_info.value.status_code == 503
+    assert exc_info.value.detail == {
+        "status": "not ready",
+        "jetstream": "JetStream not ready",
+    }
+
+
+@pytest.mark.anyio
+async def test_health_ready_timeout():
+    endorsement_processor_mock = AsyncMock(spec=EndorsementProcessor)
+
+    endorsement_processor_mock.check_jetstream.side_effect = asyncio.TimeoutError()
+
+    with pytest.raises(HTTPException) as exc_info:
+        await health_ready(endorsement_processor=endorsement_processor_mock)
+
+    assert exc_info.value.status_code == 503
+    assert exc_info.value.detail == {
+        "status": "not ready",
+        "error": "JetStream health check timed out",
+    }
+
+
+@pytest.mark.anyio
+async def test_health_ready_unexpected_error():
+    endorsement_processor_mock = AsyncMock(spec=EndorsementProcessor)
+
+    endorsement_processor_mock.check_jetstream.side_effect = Exception(
+        "Unexpected error"
+    )
+
+    with pytest.raises(HTTPException) as exc_info:
+        await health_ready(endorsement_processor=endorsement_processor_mock)
+
+    assert exc_info.value.status_code == 500
+    assert exc_info.value.detail == {"status": "error", "error": "Unexpected error"}
+
+
+@pytest.mark.anyio
+async def test_health_ready_with_timeout():
+    endorsement_processor_mock = AsyncMock(spec=EndorsementProcessor)
+
+    endorsement_processor_mock.check_jetstream.side_effect = (
+        asyncio.TimeoutError
+    )  # Simulate a slow response
+
+    with pytest.raises(HTTPException) as exc_info:
+        await health_ready(endorsement_processor=endorsement_processor_mock)
+
+    assert exc_info.value.status_code == 503
+    assert exc_info.value.detail == {
+        "status": "not ready",
+        "error": "JetStream health check timed out",
+    }