Commit e29de8ea authored by Dimitrios Gogos's avatar Dimitrios Gogos
Browse files

feat: add drain_timeout to NATS settings and enhance connection management

parent 6061d7f9

.env.example

+1 −0
Original line number Diff line number Diff line
@@ -9,3 +9,4 @@ POSTGRES_SETTINGS__CREATE_SCHEMA_ON_STARTUP = true 
NATS_SETTINGS__URL = "nats://localhost:4222"

NATS_SETTINGS__CONNECT_TIMEOUT = 10

NATS_SETTINGS__MAX_RECONNECT_ATTEMPTS = 3

NATS_SETTINGS__DRAIN_TIMEOUT = 30

.pre-commit-config.yaml

+1 −0
Original line number Diff line number Diff line
@@ -27,6 +27,7 @@ repos: 
        - "docker>=7.0.0"

        - "fastapi[standard]>=0.135.1"

        - "httpx>=0.27.0"

        - "nats-py>=2.10.0"

        - "pydantic-settings>=2.13.1"

        - "pytest>=9.0.2"

        - "pytest-asyncio>=0.24"

src/srm/adapters/databus/nats_connection_manager.py

+21 −5
Original line number Diff line number Diff line 
import asyncio



import nats

import structlog

from nats.aio.client import Client
@@ -17,6 +19,7 @@ class NatsConnectionManager: 
    def __init__(self, settings: NatsSettings) -> None:

        self._settings = settings

        self._client: Client | None = None

        self._connect_lock = asyncio.Lock()



    @property

    def is_connected(self) -> bool:
@@ -38,12 +41,17 @@ class NatsConnectionManager: 
        async def _on_reconnect() -> None:

            logger.info("nats_reconnected", url=self._settings.url)



        if not self.is_connected:

        async with self._connect_lock:

            if self.is_connected:

                logger.info("nats_already_connected")

                return



            try:

                self._client = await nats.connect(

                    servers=[self._settings.url],

                    connect_timeout=self._settings.connect_timeout,

                    max_reconnect_attempts=self._settings.max_reconnect_attempts,

                    drain_timeout=self._settings.drain_timeout,

                    error_cb=_on_error,

                    disconnected_cb=_on_disconnect,

                    reconnected_cb=_on_reconnect,
@@ -51,10 +59,18 @@ class NatsConnectionManager: 
            except Exception as e:

                logger.error("nats_error", error=str(e))

                raise

        else:

            logger.info("nats_already_connected")



    async def close(self) -> None:

        if self._client is not None:

            await self._client.drain()

        if self._client is None:

            return



        client = self._client

        try:

            # drain() is bounded by the drain_timeout passed to connect(); on timeout it

            # reports through error_cb and closes anyway, so shutdown cannot hang here.

            await client.drain()

        except Exception as e:

            logger.warning("nats_drain_failed", error=str(e))

            await client.close()

        finally:

            self._client = None

src/srm/config.py

+1 −0
Original line number Diff line number Diff line
@@ -33,6 +33,7 @@ class NatsSettings(BaseModel): 
    url: str

    connect_timeout: int

    max_reconnect_attempts: int

    drain_timeout: int = 30





class Settings(BaseSettings):

src/srm/main.py

+18 −5
Original line number Diff line number Diff line
@@ -54,11 +54,22 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: 
    except Exception as e:

        logger.error("Database engine init failed!", error=str(e))

        raise

    try:

        databus_manager: NatsConnectionManager = await init_databus_manager(

            settings=settings.nats_settings

        )

    except Exception as e:

        logger.error("Databus connection init failed!", error=str(e))

        await engine.dispose()

        raise



    try:

        databus_subscribers: list[NatsSubscriber] = await subscribe_to_subjects(databus_manager)

    except Exception as e:

        logger.error("Databus subscription failed!", error=str(e))

        await databus_manager.close()

        await engine.dispose()

        raise



    app.state.db_engine = engine

    app.state.session_maker = session_maker
@@ -68,9 +79,11 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: 
    yield



    logger.info("Shutting down application")

    await engine.dispose()

    # Drain the bus first: draining delivers in-flight messages to their handlers, and

    # those handlers still need a live DB engine.

    await app.state.databus_connection_manager.close()

    app.state.databus_subscribers = []

    await engine.dispose()





def create_app() -> FastAPI: