wslink 2.1.3__tar.gz → 2.2.1__tar.gz

{wslink-2.1.3 → wslink-2.2.1}/PKG-INFO

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

{wslink-2.1.3 → wslink-2.2.1}/setup.cfg

@@ -1,5 +1,5 @@
 [metadata]
-version = 2.1.3
+version = 2.2.1
 
 [egg_info]
 tag_build = 

{wslink-2.1.3 → wslink-2.2.1}/src/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-2.1.3 → wslink-2.2.1}/src/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-2.1.3 → wslink-2.2.1}/src/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-2.1.3 → wslink-2.2.1}/src/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,57 @@ class LinkProtocol(object):
 # =============================================================================
 
 
+class NetworkMonitor:
+    """
+    Provide context manager for increase/decrease pending request
+    either synchronously or asynchronously.
+
+    The Asynchronous version also await completion.
+    """
+
+    def __init__(self):
+
+        self.pending = 0
+        self.event = asyncio.Event()
+
+    def network_call_completed(self):
+        """Trigger completion event"""
+        self.event.set()
+
+    def on_enter(self, *args, **kwargs):
+        """Increase pending request"""
+        self.pending += 1
+
+    def on_exit(self, *args, **kwargs):
+        """Decrease pending request and trigger completion event if we reach 0 pending request"""
+        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()
+        await self.completion()
+
+    async def completion(self):
+        """Await completion of any pending network request"""
+        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 +124,7 @@ class ServerProtocol(object):
     """
 
     def __init__(self):
+        self.network_monitor = NetworkMonitor()
         self.linkProtocols = []
         self.secret = None
         self.initialize()

{wslink-2.1.3 → wslink-2.2.1}/src/wslink.egg-info/PKG-INFO

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