wslink 2.1.3__py3-none-any.whl → 2.2.0__py3-none-any.whl

wslink/__init__.py

@@ -54,7 +54,7 @@ def schedule_callback(delay, callback, *args, **kwargs):
     return loop.call_later(delay, functools.partial(callback, *args, **kwargs))
 
 
-def schedule_coroutine(delay, coro_func, *args, **kwargs):
+def schedule_coroutine(delay, coro_func, *args, done_callback=None, **kwargs):
     """
     Creates a coroutine out of the provided coroutine function coro_func and
     the provided args and kwargs, then schedules the coroutine to be called
@@ -73,4 +73,11 @@ def schedule_coroutine(delay, coro_func, *args, **kwargs):
     # See method above for comment on "get_event_loop()" vs "get_running_loop()".
     loop = asyncio.get_event_loop()
     coro_partial = functools.partial(coro_func, *args, **kwargs)
+    if done_callback is not None:
+        return loop.call_later(
+            delay,
+            lambda: asyncio.ensure_future(coro_partial()).add_done_callback(
+                done_callback
+            ),
+        )
     return loop.call_later(delay, lambda: asyncio.ensure_future(coro_partial()))

wslink/protocol.py

@@ -152,6 +152,7 @@ class WslinkHandler(object):
         self.attachment_atomic = asyncio.Lock()
         self.pub_manager = PublishManager()
         self.unchunkers = {}
+        self.network_monitor = protocol.network_monitor
 
         # Build the rpc method dictionary, assuming we were given a serverprotocol
         if self.getServerProtocol():
@@ -257,7 +258,8 @@ class WslinkHandler(object):
 
         full_message = self.unchunkers[client_id].process_chunk(msg.data)
         if full_message is not None:
-            await self.onCompleteMessage(full_message, client_id)
+            with self.network_monitor:
+                await self.onCompleteMessage(full_message, client_id)
 
     async def onCompleteMessage(self, rpc, client_id):
         logger.debug("wslink incoming msg %s", self.payloadWithSecretStripped(rpc))
@@ -282,23 +284,25 @@ class WslinkHandler(object):
 
         # Prevent any further processing if token is not valid
         if not self.isClientAuthenticated(client_id):
-            await self.sendWrappedError(
-                rpcid,
-                AUTHENTICATION_ERROR,
-                "Unauthorized: Skip message processing",
-                client_id=client_id,
-            )
+            with self.network_monitor:
+                await self.sendWrappedError(
+                    rpcid,
+                    AUTHENTICATION_ERROR,
+                    "Unauthorized: Skip message processing",
+                    client_id=client_id,
+                )
             return
 
         # No matching method found
         if not methodName in self.functionMap:
-            await self.sendWrappedError(
-                rpcid,
-                METHOD_NOT_FOUND,
-                "Unregistered method called",
-                methodName,
-                client_id=client_id,
-            )
+            with self.network_monitor:
+                await self.sendWrappedError(
+                    rpcid,
+                    METHOD_NOT_FOUND,
+                    "Unregistered method called",
+                    methodName,
+                    client_id=client_id,
+                )
             return
 
         obj, func = self.functionMap[methodName]
@@ -308,31 +312,34 @@ class WslinkHandler(object):
             self.web_app.last_active_client_id = client_id
             results = func(*args, **kwargs)
             if inspect.isawaitable(results):
-                results = await results
+                with self.network_monitor:
+                    results = await results
 
             if self.connections[client_id].closed:
                 # Connection was closed during RPC call.
                 return
 
-            await self.sendWrappedMessage(
-                rpcid, results, method=methodName, client_id=client_id
-            )
+            with self.network_monitor:
+                await self.sendWrappedMessage(
+                    rpcid, results, method=methodName, client_id=client_id
+                )
         except Exception as e_inst:
             captured_trace = traceback.format_exc()
             logger.error("Exception raised")
             logger.error(repr(e_inst))
             logger.error(captured_trace)
-            await self.sendWrappedError(
-                rpcid,
-                EXCEPTION_ERROR,
-                "Exception raised",
-                {
-                    "method": methodName,
-                    "exception": repr(e_inst),
-                    "trace": captured_trace,
-                },
-                client_id=client_id,
-            )
+            with self.network_monitor:
+                await self.sendWrappedError(
+                    rpcid,
+                    EXCEPTION_ERROR,
+                    "Exception raised",
+                    {
+                        "method": methodName,
+                        "exception": repr(e_inst),
+                        "trace": captured_trace,
+                    },
+                    client_id=client_id,
+                )
 
     def payloadWithSecretStripped(self, payload):
         payload = copy.deepcopy(payload)
@@ -397,13 +404,14 @@ class WslinkHandler(object):
         except Exception:
             # the content which is not serializable might be arbitrarily large, don't include.
             # repr(content) would do that...
-            await self.sendWrappedError(
-                rpcid,
-                RESULT_SERIALIZE_ERROR,
-                "Method result cannot be serialized",
-                method,
-                client_id=client_id,
-            )
+            with self.network_monitor:
+                await self.sendWrappedError(
+                    rpcid,
+                    RESULT_SERIALIZE_ERROR,
+                    "Method result cannot be serialized",
+                    method,
+                    client_id=client_id,
+                )
             return
 
         websockets = self.getAuthenticatedWebsockets(client_id, skip_last_active_client)
@@ -411,11 +419,15 @@ class WslinkHandler(object):
         # aiohttp can not handle pending ws.send_bytes()
         # tried with semaphore but got exception with >1
         # https://github.com/aio-libs/aiohttp/issues/2934
-        async with self.attachment_atomic:
-            for chunk in generate_chunks(packed_wrapper, MAX_MSG_SIZE):
-                for ws in websockets:
-                    if ws is not None:
-                        await ws.send_bytes(chunk)
+        with self.network_monitor:
+            async with self.attachment_atomic:
+                for chunk in generate_chunks(packed_wrapper, MAX_MSG_SIZE):
+                    for ws in websockets:
+                        if ws is not None:
+                            await ws.send_bytes(chunk)
+
+        # Network operation completed
+        self.network_monitor.network_call_completed()
 
     async def sendWrappedError(self, rpcid, code, message, data=None, client_id=None):
         wrapper = {
@@ -443,11 +455,15 @@ class WslinkHandler(object):
         # aiohttp can not handle pending ws.send_bytes()
         # tried with semaphore but got exception with >1
         # https://github.com/aio-libs/aiohttp/issues/2934
-        async with self.attachment_atomic:
-            for chunk in generate_chunks(packed_wrapper, MAX_MSG_SIZE):
-                for ws in websockets:
-                    if ws is not None:
-                        await ws.send_bytes(chunk)
+        with self.network_monitor:
+            async with self.attachment_atomic:
+                for chunk in generate_chunks(packed_wrapper, MAX_MSG_SIZE):
+                    for ws in websockets:
+                        if ws is not None:
+                            await ws.send_bytes(chunk)
+
+        # Network operation completed
+        self.network_monitor.network_call_completed()
 
     def publish(self, topic, data, client_id=None, skip_last_active_client=False):
         client_list = [client_id] if client_id else [c_id for c_id in self.connections]

wslink/publish.py

@@ -27,6 +27,7 @@ class PublishManager(object):
         for protocol in self.protocols:
             # The client is unknown - we send to any client who is subscribed to the topic
             rpcid = "publish:{0}:{1}".format(topic, self.publishCount)
+            protocol.network_monitor.on_enter()
             schedule_coroutine(
                 0,
                 protocol.sendWrappedMessage,
@@ -34,6 +35,8 @@ class PublishManager(object):
                 data,
                 client_id=client_id,
                 skip_last_active_client=skip_last_active_client,
+                # for schedule_coroutine call
+                done_callback=protocol.network_monitor.on_exit,
             )
 
 

wslink/websocket.py

@@ -5,6 +5,7 @@ ServerProtocol to hook all the needed LinkProtocols together.
 """
 
 import logging
+import asyncio
 
 from . import register as exportRpc
 from . import schedule_callback
@@ -65,6 +66,42 @@ class LinkProtocol(object):
 # =============================================================================
 
 
+class NetworkMonitor:
+    def __init__(self):
+        self.pending = 0
+        self.event = asyncio.Event()
+
+    def network_call_completed(self):
+        self.event.set()
+
+    def on_enter(self, *args, **kwargs):
+        self.pending += 1
+
+    def on_exit(self, *args, **kwargs):
+        self.pending -= 1
+        if self.pending == 0 and not self.event.is_set():
+            self.event.set()
+
+    # Sync ctx manager
+    def __enter__(self):
+        self.on_enter()
+        return self
+
+    def __exit__(self, exc_type, exc_value, exc_traceback):
+        self.on_exit()
+
+    # Async ctx manager
+    async def __aenter__(self):
+        self.on_enter()
+        return self
+
+    async def __aexit__(self, exc_t, exc_v, exc_tb):
+        self.on_exit()
+        while self.pending:
+            self.event.clear()
+            await self.event.wait()
+
+
 class ServerProtocol(object):
     """
     Defines the core server protocol for wslink. Gathers a list of LinkProtocol
@@ -72,6 +109,7 @@ class ServerProtocol(object):
     """
 
     def __init__(self):
+        self.network_monitor = NetworkMonitor()
         self.linkProtocols = []
         self.secret = None
         self.initialize()

{wslink-2.1.3.dist-info → wslink-2.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: wslink
-Version: 2.1.3
+Version: 2.2.0
 Summary: Python/JavaScript library for communicating over WebSocket
 Home-page: https://github.com/kitware/wslink
 Author: Kitware, Inc.

{wslink-2.1.3.dist-info → wslink-2.2.0.dist-info}/RECORD

@@ -1,14 +1,14 @@
 wslink/LICENSE,sha256=I44UH7kDVqxDLnnlOWw_hFL2Fz7RjQ_4vPzZv9NYgTU,1483
-wslink/__init__.py,sha256=7Uyn6lk5LhykHbBwQ9pwSmX1YhxDanD9TT6H95DXzhk,2709
+wslink/__init__.py,sha256=qdLGVhzpEXBoraTnLr5S-8bX9oK8pZCQWe8rZRlwAco,2946
 wslink/chunking.py,sha256=1DJlGG6fjknGFrqPOtqUUc5tCrijldP7Kdx56d5e3Wg,7337
 wslink/launcher.py,sha256=8VMs3juObLkyGYQFNLjMoo4qFpKIcxWz0kS-af-DKO4,21170
-wslink/protocol.py,sha256=V5kqfHebfyXzmZq0-ua0SjDkOXdQmx1Mu29G3hORgx0,15971
-wslink/publish.py,sha256=9G5TXqyGr-LCo_LwHYhzif6lhG2iXDvEBmEgwR8fh1M,1437
+wslink/protocol.py,sha256=T3F_Y_pZ4qXUApjpDoi9zXDkDvp3qnlKQICj2jnPlsA,16744
+wslink/publish.py,sha256=Xyv9piZT4HxO5l_SA7zeReH8t6tisVHUgU57Hjzgkcg,1595
 wslink/relay.py,sha256=E8Lzu2Ay7KbOheN1-ArAZawo8lLqdDgJXOZSBuMknYs,86
 wslink/server.py,sha256=yvhCjpzPOfhbZrpDvW9i4H_uSyuQAe3ZOP-BRBmgHQA,9326
 wslink/ssl_context.py,sha256=hNOJJCdrStws1Qf6vPvY4vTk9Bf8J5d90W3fS0cRv8o,2290
 wslink/uri.py,sha256=woCQ4yChUqTMg9IT6YYDtUYeKmCg7OUCEgeBGA-19DY,384
-wslink/websocket.py,sha256=pBiWqkL8Zn8LuSJ9nv3yA-KjEynbolOQ2gLHtQFJ2Ic,4611
+wslink/websocket.py,sha256=Gdp8ll5z3wq6gBU0GSVu2iq_o598Mxt8jQ1wZi6Z87A,5520
 wslink/backends/__init__.py,sha256=cyJGjm-YyBSyOEX81owyTbJ3YnrA6dB7--B4LnsEtHI,1214
 wslink/backends/aiohttp/__init__.py,sha256=Bc3uPVrhNXnn5Xt3Tck2Th-dsysz1VAnQdh8WzRn1y0,9253
 wslink/backends/aiohttp/launcher.py,sha256=gHNMvtgNHEwBN_QBRDSCrTp2B4K1PsfV81rKaHi7Cxo,8897
@@ -19,7 +19,7 @@ wslink/backends/jupyter/__init__.py,sha256=Qu65gWsd2xCSsxybnDtEDI5vMjHN-F5jgPZOy
 wslink/backends/jupyter/core.py,sha256=PCQN-uZPFROnRv8B5dNwnwHV67o4Bpme3_Z6V-zbOUA,3864
 wslink/backends/tornado/__init__.py,sha256=Qu65gWsd2xCSsxybnDtEDI5vMjHN-F5jgPZOyNIxnGs,112
 wslink/backends/tornado/core.py,sha256=tPMkkhWuO_ovkisVim0zcegwZKEAG4IRUdd_O_0a_R0,2157
-wslink-2.1.3.dist-info/METADATA,sha256=MY8aB802QkaDhIEEw7ekhqlW371QtyqLk4uws5kMqwY,3120
-wslink-2.1.3.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
-wslink-2.1.3.dist-info/top_level.txt,sha256=N0d8eqvhwhfW1p1yPTmvxlbzhjz7ZyhBfysNvaFqpQY,7
-wslink-2.1.3.dist-info/RECORD,,
+wslink-2.2.0.dist-info/METADATA,sha256=SEj4XM87pZO5pMxU6WTxYxyGz7rZdxFOspQkHEYTkEg,3120
+wslink-2.2.0.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
+wslink-2.2.0.dist-info/top_level.txt,sha256=N0d8eqvhwhfW1p1yPTmvxlbzhjz7ZyhBfysNvaFqpQY,7
+wslink-2.2.0.dist-info/RECORD,,