Repository: list-family/ansq (language: python)

Issue:
Auto-reconnect doesn't work properly
The ansq connection re-connects only once and if it fails it would stop `nsq.messages()` generator immediately - which most like cause to stop consumer service. The connection with enabled auto-reconnect feature should try to reconnect each N seconds (or progressively) and wait until the connection is established.
Run test nsqd in a subprocess
Benefits:
- easier tests run - no need to run manually nsqd, just install it and you are ready to go
- run multiple nsqd servers simultaneously
- more control over nsqd process during tests - start/stop whenever you need to
- test multiple nsq versions in CI


This PR also would help with #39 to make integrational/functioanal tests for auto-reconnect feature and start/stop the server during the tests.


Discussion / hints:

Fix patch:
diff --git a/ansq/tcp/connection.py b/ansq/tcp/connection.py
index 01c8d98..7a2bb48 100644
--- a/ansq/tcp/connection.py
+++ b/ansq/tcp/connection.py
@@ -25,6 +25,11 @@
 from ansq.typedefs import TCPResponse
 from ansq.utils import truncate_text, validate_topic_channel_name
 
+# Auto reconnect settings
+AUTO_RECONNECT_INITIAL_INTERVAL = 2
+AUTO_RECONNECT_MAX_INTERVAL = 2048
+AUTO_RECONNECT_PROGRESSION_RATIO = 2
+
 
 class NSQConnection(NSQConnectionBase):
     async def connect(self) -> bool:
@@ -59,7 +64,7 @@ async def reconnect(self, raise_error: bool = True) -> bool:
         self.logger.debug(f"Reconnecting to {self.endpoint}...")
         self._status = ConnectionStatus.RECONNECTING
 
-        await self._do_close(change_status=False)
+        await self._do_close(change_status=False, silent=True)
         try:
             await self.connect()
             await self.identify()
@@ -81,8 +86,42 @@ async def reconnect(self, raise_error: bool = True) -> bool:
         self._status = ConnectionStatus.CONNECTED
         return True
 
+    async def _do_auto_reconnect(
+        self, interval: int = AUTO_RECONNECT_INITIAL_INTERVAL
+    ) -> None:
+        """Call ``reconnect()`` method. If failed, sleep and try again."""
+        # Reconnect silently
+        try:
+            reconnected = await self.reconnect()
+        except Exception as exc:
+            await self._do_close(exc, change_status=False, silent=True)
+            reconnected = False
+
+        # Return early if succeeded
+        if reconnected:
+            return
+
+        # Don't sleep more than max interval
+        if interval > AUTO_RECONNECT_MAX_INTERVAL:
+            interval = AUTO_RECONNECT_MAX_INTERVAL
+
+        self.logger.debug(
+            "Failed to reconnect to %s. Wait for %s seconds ...",
+            self.endpoint,
+            interval,
+        )
+
+        # Reconnection is failed - sleep and schedule new reconnect
+        await asyncio.sleep(interval)
+        self._reconnect_task = self._loop.create_task(
+            self._do_auto_reconnect(interval * AUTO_RECONNECT_PROGRESSION_RATIO),
+        )
+
     async def _do_close(
-        self, exception: Optional[Exception] = None, change_status: bool = True
+        self,
+        exception: Optional[Exception] = None,
+        change_status: bool = True,
+        silent: bool = False,
     ) -> None:
         if self.is_closed or self.status.is_init or self.status.is_closing:
             return
@@ -90,14 +129,15 @@ async def _do_close(
         if change_status:
             self._status = ConnectionStatus.CLOSING
 
-        if exception:
-            self.logger.error(
-                "Connection {} is closing due an error: {}".format(
-                    self.endpoint, exception,
-                ),
-            )
-        else:
-            self.logger.debug(f"Connection {self.endpoint} is closing...")
+        if not silent:
+            if exception:
+                self.logger.error(
+                    "Connection {} is closing due an error: {}".format(
+                        self.endpoint, exception,
+                    ),
+                )
+            else:
+                self.logger.debug(f"Connection {self.endpoint} is closing...")
 
         if self.is_subscribed and change_status:
             self._is_subscribed = False
@@ -274,9 +314,7 @@ async def _read_data_task(self) -> None:
         self.logger.info("Lost connection to NSQ")
         if self._auto_reconnect:
             await asyncio.sleep(1)
-            self._reconnect_task = self._loop.create_task(
-                self.reconnect(raise_error=False),
-            )
+            self._reconnect_task = self._loop.create_task(self._do_auto_reconnect())
         else:
             await self._do_close(OSError("Lost connection to NSQ"))
 
diff --git a/ansq/tcp/types/connection.py b/ansq/tcp/types/connection.py
index f646af1..63e50f4 100644
--- a/ansq/tcp/types/connection.py
+++ b/ansq/tcp/types/connection.py
@@ -165,7 +165,10 @@ async def cls(self) -> None:
 
     @abc.abstractmethod
     async def _do_close(
-        self, exception: Optional[Exception] = None, change_status: bool = True
+        self,
+        exception: Optional[Exception] = None,
+        change_status: bool = True,
+        silent: bool = False,
     ) -> None:
         raise NotImplementedError()
 


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index 70f7a3a..3fa70e3 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,9 +1,11 @@
 import asyncio
+import inspect
 import os
 import shutil
 import signal
 import time
 from asyncio.subprocess import Process
+from typing import Awaitable, Callable, Union
 
 import async_generator
 import pytest
@@ -114,3 +116,33 @@ async def _create_nsqd(host="127.0.0.1", port=4150, http_port=4151):
 async def nsqd(create_nsqd) -> AsyncNSQD:
     async with create_nsqd() as nsqd:
         yield nsqd
+
+
+@pytest.fixture
+def wait_for():
+    """ Wait for a predicate with a timeout."""
+
+    async def inner(
+        predicate: Union[Callable[..., bool], Callable[..., Awaitable[bool]]],
+        timeout: float = 5.0,
+        sleep_time: float = 0.1,
+    ):
+        __tracebackhide__ = True
+
+        start = time.time()
+
+        while True:
+            predicate_result = (
+                await predicate()
+                if inspect.iscoroutinefunction(predicate)
+                else predicate()
+            )
+            if predicate_result:
+                return
+
+            if time.time() - start > timeout:  # pragma: no cover
+                raise AssertionError("failed to wait for predicate")
+
+            await asyncio.sleep(sleep_time)
+
+    return inner
diff --git a/tests/test_connection.py b/tests/test_connection.py
index 882964e..0132e87 100644
--- a/tests/test_connection.py
+++ b/tests/test_connection.py
@@ -37,3 +37,18 @@ async def test_reconnect_while_connected():
 
     await nsq.close()
     assert nsq.status.is_closed
+
+
+@pytest.mark.asyncio
+async def test_auto_reconnect(nsqd, wait_for):
+    nsq = await open_connection(auto_reconnect=True)
+    assert nsq.status.is_connected
+
+    await nsqd.stop()
+    await wait_for(lambda: nsq.status.is_reconnecting)
+
+    await nsqd.start()
+    await wait_for(lambda: nsq.status.is_connected)
+
+    await nsq.close()
+    assert nsq.status.is_closed