added test for websocket reconnect flow.

better reconnect flow

Author: [email protected]

rsocket/rsocket_base.py

@@ -79,6 +79,7 @@ def __init__(self,
         self._responder_lease = None
         self._requester_lease = None
         self._is_closing = False
+        self._connecting = True
 
         self._async_frame_handler_by_type: Dict[Type[Frame], Any] = {
             RequestResponseFrame: self.handle_request_response,
@@ -420,7 +421,7 @@ async def close(self):
     async def _close_transport(self):
         if self._current_transport().done():
             logger().debug('%s: Closing transport', self._log_identifier())
-            transport = self._current_transport().result()
+            transport = await self._current_transport()
 
             if transport is not None:
                 try:

rsocket/rsocket_client.py

@@ -70,14 +70,17 @@ async def connect(self):
         return await super().connect()
 
     async def _connect_new_transport(self):
-        new_transport = await self._get_new_transport()
+        try:
+            new_transport = await self._get_new_transport()
 
-        if new_transport is None:
-            raise RSocketNoAvailableTransport()
+            if new_transport is None:
+                raise RSocketNoAvailableTransport()
 
-        self._next_transport.set_result(new_transport)
-        transport = await self._current_transport()
-        await transport.connect()
+            self._next_transport.set_result(new_transport)
+            transport = await self._current_transport()
+            await transport.connect()
+        finally:
+            self._connecting = False
 
     async def _get_new_transport(self):
         try:
@@ -110,6 +113,13 @@ async def _reconnect_listener(self):
         try:
             while True:
                 await self._connect_request_event.wait()
+
+                logger().debug('%s: Got reconnect request', self._log_identifier())
+
+                if self._connecting:
+                    continue
+
+                self._connecting = True
                 self._connect_request_event.clear()
                 await self._close(reconnect=True)
                 self._next_transport = create_future()

rsocket/transports/abstract_messaging.py

@@ -12,6 +12,9 @@ def __init__(self):
     async def next_frame_generator(self):
         frame = await self._incoming_frame_queue.get()
 
+        if isinstance(frame, Exception):
+            raise frame
+
         async def frame_generator():
             yield frame
 

rsocket/transports/aiohttp_websocket.py

@@ -4,6 +4,7 @@
 import aiohttp
 from aiohttp import web
 
+from rsocket.exceptions import RSocketTransportError
 from rsocket.frame import Frame
 from rsocket.helpers import wrap_transport_exception, single_transport_provider
 from rsocket.logger import logger
@@ -44,20 +45,29 @@ def __init__(self, url):
         self._ws_context = None
         self._ws = None
         self._message_handler = None
+        self._connection_ready = asyncio.Event()
 
     async def connect(self):
         self._session = aiohttp.ClientSession()
         self._ws_context = self._session.ws_connect(self._url)
         self._ws = await self._ws_context.__aenter__()
+        self._connection_ready.set()
         self._message_handler = asyncio.create_task(self.handle_incoming_ws_messages())
 
     async def handle_incoming_ws_messages(self):
-        async for msg in self._ws:
-            if msg.type == aiohttp.WSMsgType.BINARY:
-                async for frame in self._frame_parser.receive_data(msg.data, 0):
-                    self._incoming_frame_queue.put_nowait(frame)
+        await self._connection_ready.wait()
+        try:
+            async for msg in self._ws:
+                if msg.type == aiohttp.WSMsgType.BINARY:
+                    async for frame in self._frame_parser.receive_data(msg.data, 0):
+                        self._incoming_frame_queue.put_nowait(frame)
+        except asyncio.CancelledError:
+            pass
+        except Exception:
+            self._incoming_frame_queue.put_nowait(RSocketTransportError())
 
     async def send_frame(self, frame: Frame):
+        await self._connection_ready.wait()
         with wrap_transport_exception():
             await self._ws.send_bytes(frame.serialize())
 

tests/rsocket/helpers.py

@@ -1,4 +1,5 @@
 import asyncio
+from dataclasses import dataclass
 from datetime import timedelta
 from math import ceil
 from typing import Type
@@ -8,6 +9,8 @@
 from rsocket.payload import Payload
 from rsocket.request_handler import BaseRequestHandler
 from rsocket.rsocket_base import RSocketBase
+from rsocket.rsocket_server import RSocketServer
+from rsocket.transports.transport import Transport
 
 
 def data_bits(data: bytes, name: str = None):
@@ -69,6 +72,12 @@ def factory(self, socket) -> BaseRequestHandler:
         return self._handler_factory(socket, self._server_id, self._delay)
 
 
-async def force_closing_connection(current_connection, delay=timedelta(0)):
+async def force_closing_connection(transport, delay=timedelta(0)):
     await asyncio.sleep(delay.total_seconds())
-    current_connection[1].close()
+    await transport.close()
+
+
+@dataclass
+class ServerContainer:
+    server: RSocketServer = None
+    transport: Transport = None

tests/rsocket/test_connection_lost.py

@@ -5,6 +5,8 @@
 from typing import Optional, Tuple
 
 import pytest
+from aiohttp.test_utils import RawTestServer
+from asyncstdlib import sync
 
 from reactivestreams.publisher import Publisher
 from rsocket.awaitable.awaitable_rsocket import AwaitableRSocket
@@ -17,10 +19,11 @@
 from rsocket.rsocket_client import RSocketClient
 from rsocket.rsocket_server import RSocketServer
 from rsocket.streams.stream_from_async_generator import StreamFromAsyncGenerator
+from rsocket.transports.aiohttp_websocket import websocket_handler_factory, TransportAioHttpClient
 from rsocket.transports.tcp import TransportTCP
 from rsocket.transports.transport import Transport
 from tests.rsocket.helpers import future_from_payload, IdentifiedHandlerFactory, \
-    IdentifiedHandler, force_closing_connection
+    IdentifiedHandler, force_closing_connection, ServerContainer
 
 
 class ServerHandler(IdentifiedHandler):
@@ -45,7 +48,7 @@ async def test_connection_lost(unused_tcp_port):
     index_iterator = iter(range(1, 3))
 
     wait_for_server = Event()
-    server_connection: Optional[Tuple] = None
+    transport: Optional[Transport] = None
     client_connection: Optional[Tuple] = None
 
     class ClientHandler(BaseRequestHandler):
@@ -54,9 +57,9 @@ async def on_connection_lost(self, rsocket, exception: Exception):
             await rsocket.reconnect()
 
     def session(*connection):
-        nonlocal server, server_connection
-        server_connection = connection
-        server = RSocketServer(TransportTCP(*connection),
+        nonlocal server, transport
+        transport = TransportTCP(*connection)
+        server = RSocketServer(transport,
                                IdentifiedHandlerFactory(next(index_iterator), ServerHandler).factory)
         wait_for_server.set()
 
@@ -89,7 +92,9 @@ async def transport_provider():
             await wait_for_server.wait()
             wait_for_server.clear()
             response1 = await connection.request_response(Payload(b'request 1'))
-            await force_closing_connection(server_connection)
+
+            await force_closing_connection(transport)
+
             await server.close()  # cleanup async tasks from previous server to avoid errors (?)
             await wait_for_server.wait()
             response2 = await connection.request_response(Payload(b'request 2'))
@@ -118,11 +123,11 @@ async def close(self):
 
 
 @pytest.mark.allow_error_log(regex_filter='Connection error')
-async def test_connection_failure(unused_tcp_port: int):
+async def test_tcp_connection_failure(unused_tcp_port: int):
     index_iterator = iter(range(1, 3))
 
     wait_for_server = Event()
-    server_connection: Optional[Tuple] = None
+    transport: Optional[Transport] = None
     client_connection: Optional[Tuple] = None
 
     class ClientHandler(BaseRequestHandler):
@@ -131,9 +136,9 @@ async def on_connection_lost(self, rsocket, exception: Exception):
             await rsocket.reconnect()
 
     def session(*connection):
-        nonlocal server, server_connection
-        server_connection = connection
-        server = RSocketServer(TransportTCP(*connection),
+        nonlocal server, transport
+        transport = TransportTCP(*connection)
+        server = RSocketServer(transport,
                                IdentifiedHandlerFactory(next(index_iterator), ServerHandler).factory)
         wait_for_server.set()
 
@@ -172,7 +177,7 @@ async def transport_provider():
 
             response1 = await connection.request_response(Payload(b'request 1'))
 
-            await force_closing_connection(server_connection)
+            await force_closing_connection(transport)
 
             await server.close()  # cleanup async tasks from previous server to avoid errors (?)
             await wait_for_server.wait()
@@ -187,75 +192,115 @@ async def transport_provider():
         service.close()
 
 
-@pytest.mark.allow_error_log(regex_filter='Connection error')
-async def test_connection_failure_during_stream(unused_tcp_port):
-    index_iterator = iter(range(1, 3))
+class ClientHandler(BaseRequestHandler):
+    async def on_connection_lost(self, rsocket, exception: Exception):
+        logger().info('Test Reconnecting')
+        await rsocket.reconnect()
 
-    wait_for_server = Event()
-    server_connection: Optional[Tuple] = None
-    client_connection: Optional[Tuple] = None
 
-    class ClientHandler(BaseRequestHandler):
-        async def on_connection_lost(self, rsocket, exception: Exception):
-            logger().info('Test Reconnecting')
-            await rsocket.reconnect()
+async def start_tcp_service(waiter: asyncio.Event, container, port: int):
+    index_iterator = iter(range(1, 3))
 
     def session(*connection):
-        nonlocal server, server_connection
-        server_connection = connection
-        server = RSocketServer(TransportTCP(*connection),
-                               IdentifiedHandlerFactory(next(index_iterator),
-                                                        ServerHandler,
-                                                        delay=timedelta(seconds=1)).factory)
-        wait_for_server.set()
+        container.transport = TransportTCP(*connection)
+        container.server = RSocketServer(container.transport,
+                                         IdentifiedHandlerFactory(next(index_iterator),
+                                                                  ServerHandler,
+                                                                  delay=timedelta(seconds=1)).factory)
+        waiter.set()
 
-    async def start():
-        nonlocal service, client
-        service = await asyncio.start_server(session, host, port)
+    service = await asyncio.start_server(session, 'localhost', port)
+    return sync(service.close)
 
-        async def transport_provider():
-            try:
-                nonlocal client_connection
-                client_connection = await asyncio.open_connection(host, port)
-                yield TransportTCP(*client_connection)
 
-                yield FailingTransportTCP()
+async def start_tcp_client(port: int) -> RSocketClient:
+    async def transport_provider():
+        try:
+            client_connection = await asyncio.open_connection('localhost', port)
+            yield TransportTCP(*client_connection)
 
-                client_connection = await asyncio.open_connection(host, port)
-                yield TransportTCP(*client_connection)
-            except Exception:
-                logger().error('Client connection error', exc_info=True)
-                raise
+            yield FailingTransportTCP()
 
-        client = RSocketClient(transport_provider(), handler_factory=ClientHandler)
+            client_connection = await asyncio.open_connection('localhost', port)
+            yield TransportTCP(*client_connection)
+        except Exception:
+            logger().error('Client connection error', exc_info=True)
+            raise
 
-    service: Optional[Server] = None
-    server: Optional[RSocketServer] = None
-    client: Optional[RSocketClient] = None
-    port = unused_tcp_port
-    host = 'localhost'
+    return RSocketClient(transport_provider(), handler_factory=ClientHandler)
 
-    await start()
+
+async def start_websocket_service(waiter: asyncio.Event, container, port: int):
+    index_iterator = iter(range(1, 3))
+
+    def handler_factory(*args, **kwargs):
+        return IdentifiedHandlerFactory(
+            next(index_iterator),
+            ServerHandler,
+            delay=timedelta(seconds=1)).factory(*args, **kwargs)
+
+    def on_server_create(server):
+        container.server = server
+        container.transport = server._transport
+        waiter.set()
+
+    server = RawTestServer(websocket_handler_factory(on_server_create=on_server_create,
+                                                     handler_factory=handler_factory), port=port)
+    await server.start_server()
+    return server.close
+
+
+async def start_websocket_client(port: int) -> RSocketClient:
+    url = 'http://localhost:{}'.format(port)
+
+    async def transport_provider():
+        try:
+            yield TransportAioHttpClient(url)
+
+            yield FailingTransportTCP()
+
+            yield TransportAioHttpClient(url)
+        except Exception:
+            logger().error('Client connection error', exc_info=True)
+            raise
+
+    return RSocketClient(transport_provider(), handler_factory=ClientHandler)
+
+
+@pytest.mark.allow_error_log(regex_filter='Connection error')
+@pytest.mark.parametrize(
+    'start_service, start_client',
+    (
+            (start_tcp_service, start_tcp_client),
+            (start_websocket_service, start_websocket_client),
+    )
+)
+async def test_connection_failure_during_stream(unused_tcp_port, start_service, start_client):
+    server_container = ServerContainer()
+    wait_for_server = Event()
+
+    service_closer = await start_service(wait_for_server, server_container, unused_tcp_port)
+    client = await start_client(unused_tcp_port)
 
     try:
-        async with AwaitableRSocket(client) as connection:
+        async with AwaitableRSocket(client) as a_client:
             await wait_for_server.wait()
             wait_for_server.clear()
 
             with pytest.raises(RSocketProtocolError) as exc_info:
                 await asyncio.gather(
-                    connection.request_stream(Payload(b'request 1')),
-                    force_closing_connection(server_connection, timedelta(seconds=2)))
+                    a_client.request_stream(Payload(b'request 1')),
+                    force_closing_connection(server_container.transport, timedelta(seconds=2)))
 
             assert exc_info.value.data == 'Connection error'
             assert exc_info.value.error_code == ErrorCode.CONNECTION_ERROR
 
-            await server.close()  # cleanup async tasks from previous server to avoid errors (?)
+            await server_container.server.close()  # cleanup async tasks from previous server to avoid errors (?)
             await wait_for_server.wait()
-            response2 = await connection.request_response(Payload(b'request 2'))
+            response2 = await a_client.request_response(Payload(b'request 2'))
 
             assert response2.data == b'data: request 2 server 2'
     finally:
-        await server.close()
+        await server_container.server.close()
 
-        service.close()
+        await service_closer()