[Pyamqp] Async WS implementation using a pure async library

sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_connection_async.py

@@ -185,7 +185,7 @@ async def _disconnect(self, *args) -> None:
         if self.state == ConnectionState.END:
             return
         await self._set_state(ConnectionState.END)
-        self._transport.close()
+        await self._transport.close()
 
     def _can_read(self):
         # type: () -> bool

sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_transport_async.py

@@ -37,9 +37,12 @@
 import socket
 import ssl
 import struct
-from ssl import SSLError
+from ssl import SSLContext, SSLError
+from contextlib import contextmanager
 from io import BytesIO
 import logging
+from threading import Lock
+from typing import Optional
 
 import certifi
 
@@ -57,6 +60,7 @@
     set_cloexec,
     AMQP_PORT,
     TIMEOUT_INTERVAL,
+    WebSocketTransport,
 )
 
 
@@ -145,37 +149,6 @@ async def send_frame(self, channel, frame, **kwargs):
         # _LOGGER.info("OCH%d -> %r", channel, frame)
 
 
-class AsyncTransport(AsyncTransportMixin):
-    """Common superclass for TCP and SSL transports."""
-
-    def __init__(
-        self,
-        host,
-        port=AMQP_PORT,
-        connect_timeout=None,
-        read_timeout=None,
-        write_timeout=None,
-        ssl=False,
-        socket_settings=None,
-        raise_on_initial_eintr=True,
-        **kwargs
-    ):
-        self.connected = False
-        self.sock = None
-        self.reader = None
-        self.writer = None
-        self.raise_on_initial_eintr = raise_on_initial_eintr
-        self._read_buffer = BytesIO()
-        self.host, self.port = to_host_port(host, port)
-
-        self.connect_timeout = connect_timeout
-        self.read_timeout = read_timeout
-        self.write_timeout = write_timeout
-        self.socket_settings = socket_settings
-        self.loop = get_running_loop()
-        self.socket_lock = asyncio.Lock()
-        self.sslopts = self._build_ssl_opts(ssl)
-
     def _build_ssl_opts(self, sslopts):
         if sslopts in [True, False, None, {}]:
             return sslopts
@@ -214,6 +187,40 @@ def _build_ssl_context(self, sslopts, check_hostname=None, **ctx_options):
         ctx.check_hostname = check_hostname
         return ctx
 
+
+class AsyncTransport(AsyncTransportMixin):
+    """Common superclass for TCP and SSL transports."""
+
+    def __init__(
+        self,
+        host,
+        port=AMQP_PORT,
+        connect_timeout=None,
+        read_timeout=None,
+        write_timeout=None,
+        ssl=False,
+        socket_settings=None,
+        raise_on_initial_eintr=True,
+        **kwargs,
+    ):
+        self.connected = False
+        self.sock = None
+        self.reader = None
+        self.writer = None
+        self.raise_on_initial_eintr = raise_on_initial_eintr
+        self._read_buffer = BytesIO()
+        self.host, self.port = to_host_port(host, port)
+
+        self.connect_timeout = connect_timeout
+        self.read_timeout = read_timeout
+        self.write_timeout = write_timeout
+        self.socket_settings = socket_settings
+        self.loop = get_running_loop()
+        self.socket_lock = asyncio.Lock()
+        self.sslopts = self._build_ssl_opts(ssl)
+
+
+
     async def connect(self):
         try:
             # are we already connected?
@@ -263,8 +270,7 @@ async def _connect(self, host, port, timeout):
                     # if getaddrinfo succeeded before for another address
                     # family, reraise the previous socket.error since it's more
                     # relevant to users
-                    raise e if e is not None else socket.error("failed to resolve broker hostname")
-                continue  # pragma: no cover
+                    raise (e if e is not None else socket.error("failed to resolve broker hostname"))
 
             # now that we have address(es) for the hostname, connect to broker
             for i, res in enumerate(entries):
@@ -378,7 +384,7 @@ async def _write(self, s):
         """Write a string out to the SSL socket fully."""
         self.writer.write(s)
 
-    def close(self):
+    async def close(self):
         if self.writer is not None:
             if self.sslopts:
                 # see issue: https://github.com/encode/httpx/issues/914
@@ -426,74 +432,81 @@ async def negotiate(self):
 class WebSocketTransportAsync(AsyncTransportMixin):
     def __init__(self, host, port=WEBSOCKET_PORT, connect_timeout=None, ssl=None, **kwargs):
         self._read_buffer = BytesIO()
-        self.loop = get_running_loop()
         self.socket_lock = asyncio.Lock()
-        self.sslopts = ssl if isinstance(ssl, dict) else {}
+        self.sslopts = self._build_ssl_opts(ssl) if isinstance(ssl, dict) else None
         self._connect_timeout = connect_timeout or TIMEOUT_INTERVAL
         self._custom_endpoint = kwargs.get("custom_endpoint")
         self.host = host
         self.ws = None
+        self.session = None
         self._http_proxy = kwargs.get("http_proxy", None)
-
     async def connect(self):
-        http_proxy_host, http_proxy_port, http_proxy_auth = None, None, None
+        username, password = None, None
+        http_proxy_host, http_proxy_port = None, None
+        http_proxy_auth = None
+
         if self._http_proxy:
             http_proxy_host = self._http_proxy["proxy_hostname"]
             http_proxy_port = self._http_proxy["proxy_port"]
+            if http_proxy_host and http_proxy_port:
+                http_proxy_host = f"{http_proxy_host}:{http_proxy_port}"
             username = self._http_proxy.get("username", None)
             password = self._http_proxy.get("password", None)
-            if username or password:
-                http_proxy_auth = (username, password)
+
         try:
-            from websocket import create_connection
+            from aiohttp import ClientSession
 
-            self.ws = create_connection(
+            if username or password:
+                from aiohttp import BasicAuth
+                http_proxy_auth = BasicAuth(login=username, password=password)
+
+            self.session = ClientSession()
+            self.ws = await self.session.ws_connect(
                 url="wss://{}".format(self._custom_endpoint or self.host),
-                subprotocols=[AMQP_WS_SUBPROTOCOL],
                 timeout=self._connect_timeout,
-                skip_utf8_validation=True,
-                sslopt=self.sslopts,
-                http_proxy_host=http_proxy_host,
-                http_proxy_port=http_proxy_port,
-                http_proxy_auth=http_proxy_auth,
+                protocols=[AMQP_WS_SUBPROTOCOL],
+                autoclose=False,
+                proxy=http_proxy_host,
+                proxy_auth=http_proxy_auth,
+                ssl=self.sslopts,
             )
+
         except ImportError:
-            raise ValueError("Please install websocket-client library to use websocket transport.")
+            raise ValueError("Please install aiohttp library to use websocket transport.")
 
     async def _read(self, n, buffer=None, **kwargs):  # pylint: disable=unused-arguments
         """Read exactly n bytes from the peer."""
-        from websocket import WebSocketTimeoutException
 
         length = 0
         view = buffer or memoryview(bytearray(n))
         nbytes = self._read_buffer.readinto(view)
         length += nbytes
         n -= nbytes
+
         try:
             while n:
-                data = await self.loop.run_in_executor(None, self.ws.recv)
-
+                data = await self.ws.receive_bytes()
                 if len(data) <= n:
                     view[length : length + len(data)] = data
                     n -= len(data)
                 else:
                     view[length : length + n] = data[0:n]
                     self._read_buffer = BytesIO(data[n:])
-                    n = 0
-
+                    n = 0 
             return view
-        except WebSocketTimeoutException:
+        except (asyncio.TimeoutError) as wex:
             raise TimeoutError()
 
-    def close(self):
+    async def close(self):
         """Do any preliminary work in shutting down the connection."""
-        self.ws.close()
+        await self.ws.close()
+        await self.session.close()
         self.connected = False
 
     async def write(self, s):
-        """Completely write a string to the peer.
+        """Completely write a string (byte array) to the peer.
         ABNF, OPCODE_BINARY = 0x2
         See http://tools.ietf.org/html/rfc5234
         http://tools.ietf.org/html/rfc6455#section-5.2
-        """
-        await self.loop.run_in_executor(None, self.ws.send_binary, s)
+        """       
+        await self.ws.send_bytes(s)