Switched to ws4py from websocket-client

It's ~40% faster at processing trace file messages on a pi

Author: pmeenan

Dockerfile

@@ -38,7 +38,6 @@ RUN pip install \
     psutil \
     requests \
     ujson \
-    websocket-client \
     xvfbwrapper
 
 COPY wptagent.py /wptagent/wptagent.py

docs/install.md

@@ -10,14 +10,13 @@ wptagent currently supports Windows and Linux hosts (possibly OSX but not tested
     * pypiwin32 (Windows only)
     * requests
     * ujson
-    * websocket-client
     * xvfbwrapper (Linux only)
 * Imagemagick installed and available in the path
     * The legacy tools (convert, compare, etc) need to be installed which may be optional on Windows
 * ffmpeg installed and available in the path
 * Xvfb (Linux only)
 * Debian:
-    * ```sudo apt-get install -y python2.7 python-pip imagemagick ffmpeg xvfb && sudo pip install dnspython monotonic pillow psutil requests ujson websocket-client xvfbwrapper```
+    * ```sudo apt-get install -y python2.7 python-pip imagemagick ffmpeg xvfb && sudo pip install dnspython monotonic pillow psutil requests ujson xvfbwrapper```
 * Chrome Browser
     * Linux stable channel on Ubuntu/Debian:
         * ```wget -q -O - https://dl-ssl.google.com/linux/linux_signing_key.pub | sudo apt-key add -```

internal/devtools.py

@@ -6,16 +6,17 @@
 import gzip
 import logging
 import os
-import re
+import Queue
 import subprocess
 import time
 import monotonic
 import ujson as json
+from ws4py.client.threadedclient import WebSocketClient
 
 class DevTools(object):
     """Interface into Chrome's remote dev tools protocol"""
     def __init__(self, job, task, use_devtools_video):
-        self.url = "http://localhost:{0:d}/json".format(task['port'])
+        self.url = "http://127.0.0.1:{0:d}/json".format(task['port'])
         self.websocket = None
         self.job = job
         self.task = task
@@ -84,11 +85,13 @@ def connect(self, timeout):
                                     # Close extra tabs
                                     requests.get(self.url + '/close/' + tabs[index]['id'])
                         if websocket_url is not None:
-                            from websocket import create_connection
-                            self.websocket = create_connection(websocket_url)
-                            if self.websocket:
-                                self.websocket.settimeout(1)
+                            try:
+                                self.websocket = DevToolsClient(websocket_url)
+                                self.websocket.connect()
                                 ret = True
+                            except Exception:
+                                logging.critical("Connect to dev tools websocket Error: %s",
+                                                 err.__str__())
                         else:
                             time.sleep(0.5)
                     else:
@@ -182,10 +185,9 @@ def collect_trace(self):
                 logging.info('Collecting trace events')
                 done = False
                 last_message = monotonic.monotonic()
-                self.websocket.settimeout(1)
                 while not done and monotonic.monotonic() - last_message < 30:
                     try:
-                        raw = self.websocket.recv()
+                        raw = self.websocket.get_message(1)
                         if raw is not None and len(raw):
                             msg = json.loads(raw)
                             if 'method' in msg:
@@ -346,9 +348,8 @@ def flush_pending_messages(self):
         """Clear out any pending websocket messages"""
         if self.websocket:
             try:
-                self.websocket.settimeout(0)
                 while True:
-                    raw = self.websocket.recv()
+                    raw = self.websocket.get_message(0)
                     if raw is not None and len(raw):
                         logging.debug(raw[:1000])
                         msg = json.loads(raw)
@@ -369,11 +370,10 @@ def send_command(self, method, params, wait=False, timeout=30):
                 logging.debug("Sending: %s", out)
                 self.websocket.send(out)
                 if wait:
-                    self.websocket.settimeout(1)
                     end_time = monotonic.monotonic() + timeout
                     while ret is None and monotonic.monotonic() < end_time:
                         try:
-                            raw = self.websocket.recv()
+                            raw = self.websocket.get_message(1)
                             if raw is not None and len(raw):
                                 logging.debug(raw[:1000])
                                 msg = json.loads(raw)
@@ -389,13 +389,12 @@ def send_command(self, method, params, wait=False, timeout=30):
     def wait_for_page_load(self):
         """Wait for the page load and activity to finish"""
         if self.websocket:
-            self.websocket.settimeout(1)
             start_time = monotonic.monotonic()
             end_time = start_time + self.task['time_limit']
             done = False
             while not done:
                 try:
-                    raw = self.websocket.recv()
+                    raw = self.websocket.get_message(1)
                     if raw is not None and len(raw):
                         logging.debug(raw[:1000])
                         msg = json.loads(raw)
@@ -671,3 +670,40 @@ def get_header_value(self, headers, name):
                         value = headers[header_name]
                         break
         return value
+
+class DevToolsClient(WebSocketClient):
+    """DevTools Websocket client"""
+    def __init__(self, url, protocols=None, extensions=None, heartbeat_freq=None,
+                 ssl_options=None, headers=None):
+        WebSocketClient.__init__(self, url, protocols, extensions, heartbeat_freq,
+                                 ssl_options, headers)
+        self.connected = False
+        self.messages = Queue.Queue()
+
+    def opened(self):
+        """Websocket interface - connection opened"""
+        logging.debug("DevTools websocket connected")
+        self.connected = True
+
+    def closed(self, code, reason=None):
+        """Websocket interface - connection closed"""
+        logging.debug("DevTools websocket disconnected")
+        self.connected = False
+
+    def received_message(self, raw):
+        """Websocket interface - message received"""
+        if raw.is_text:
+            message = raw.data.decode(raw.encoding) if raw.encoding is not None else raw.data
+            self.messages.put(message)
+
+    def get_message(self, timeout):
+        """Wait for and return a message from the queue"""
+        message = None
+        try:
+            if timeout is None or timeout <= 0:
+                message = self.messages.get_nowait()
+            else:
+                message = self.messages.get(True, timeout)
+        except Exception:
+            pass
+        return message

ubuntu_install.sh

@@ -1,6 +1,6 @@
 #!/bin/bash
 sudo apt-get install -y python-pip imagemagick ffmpeg xvfb
-sudo pip install dnspython monotonic pillow psutil requests ujson websocket-client xvfbwrapper
+sudo pip install dnspython monotonic pillow psutil requests ujson xvfbwrapper
 wget -q -O - https://dl-ssl.google.com/linux/linux_signing_key.pub | sudo apt-key add -
 sudo sh -c 'echo "deb http://dl.google.com/linux/chrome/deb/ stable main" >> /etc/apt/sources.list.d/google.list' 
 sudo apt-get update

wptagent.py

@@ -167,12 +167,6 @@ def startup(self):
             print "Missing ujson parser. Please run 'pip install ujson'"
             ret = False
 
-        try:
-            import websocket as _
-        except ImportError:
-            print "Missing websocket module. Please run 'pip install websocket-client'"
-            ret = False
-
         try:
             subprocess.check_output(['python', '--version'])
         except Exception:

ws4py/LICENSE

@@ -0,0 +1,29 @@
+NOTICE: The version of ws4py included with kaithem has been modified.
+
+
+Copyright (c) 2011-2015, Sylvain Hellegouarch 
+All rights reserved. 
+
+Redistribution and use in source and binary forms, with or without 
+modification, are permitted provided that the following conditions are met: 
+
+ * Redistributions of source code must retain the above copyright notice, 
+   this list of conditions and the following disclaimer. 
+ * Redistributions in binary form must reproduce the above copyright 
+   notice, this list of conditions and the following disclaimer in the 
+   documentation and/or other materials provided with the distribution. 
+ * Neither the name of ws4py nor the names of its contributors may be used 
+   to endorse or promote products derived from this software without 
+   specific prior written permission. 
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE 
+ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE 
+LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR 
+CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF 
+SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS 
+INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN 
+CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) 
+ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE 
+POSSIBILITY OF SUCH DAMAGE.

ws4py/__init__.py

@@ -0,0 +1,67 @@
+# -*- coding: utf-8 -*-
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are
+# met:
+#
+#     * Redistributions of source code must retain the above copyright
+# notice, this list of conditions and the following disclaimer.
+#     * Redistributions in binary form must reproduce the above
+# copyright notice, this list of conditions and the following disclaimer
+# in the documentation and/or other materials provided with the
+# distribution.
+#     * Neither the name of ws4py nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+import logging
+import logging.handlers as handlers
+
+__author__ = "Sylvain Hellegouarch"
+__version__ = "0.4.2.dev0"
+__all__ = ['WS_KEY', 'WS_VERSION', 'configure_logger', 'format_addresses']
+
+WS_KEY = b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
+WS_VERSION = (8, 13)
+
+def configure_logger(stdout=True, filepath=None, level=logging.INFO):
+    logger = logging.getLogger('ws4py')
+    logger.setLevel(level)
+    logfmt = logging.Formatter("[%(asctime)s] %(levelname)s %(message)s")
+
+    if filepath:
+        h = handlers.RotatingFileHandler(filepath, maxBytes=10485760, backupCount=3)
+        h.setLevel(level)
+        h.setFormatter(logfmt)
+        logger.addHandler(h)
+
+    if stdout:
+        import sys
+        h = logging.StreamHandler(sys.stdout)
+        h.setLevel(level)
+        h.setFormatter(logfmt)
+        logger.addHandler(h)
+
+    return logger
+
+def format_addresses(ws):
+    me = ws.local_address
+    peer = ws.peer_address
+    if isinstance(me, tuple) and isinstance(peer, tuple):
+        me_ip, me_port = ws.local_address
+        peer_ip, peer_port = ws.peer_address
+        return "[Local => %s:%d | Remote => %s:%d]" % (me_ip, me_port, peer_ip, peer_port)
+
+    return "[Bound to '%s']" % me

ws4py/async_websocket.py

@@ -0,0 +1,126 @@
+# -*- coding: utf-8 -*-
+__doc__ = """
+WebSocket implementation that relies on two new Python
+features:
+
+* asyncio to provide the high-level interface above transports
+* yield from to delegate to the reading stream whenever more
+  bytes are required
+
+You can use these implementations in that context
+and benefit from those features whilst using ws4py.
+
+Strictly speaking this module probably doesn't have to
+be called async_websocket but it feels this will be its typical
+usage and is probably more readable than
+delegated_generator_websocket_on_top_of_asyncio.py
+"""
+import asyncio
+import types
+
+from ws4py.websocket import WebSocket as _WebSocket
+from ws4py.messaging import Message
+
+__all__ = ['WebSocket', 'EchoWebSocket']
+
+class WebSocket(_WebSocket):
+    def __init__(self, proto):
+        """
+        A :pep:`3156` ready websocket handler that works
+        well in a coroutine-aware loop such as the one provided
+        by the asyncio module.
+
+        The provided `proto` instance is a
+        :class:`asyncio.Protocol` subclass instance that will
+        be used internally to read and write from the
+        underlying transport.
+
+        Because the base :class:`ws4py.websocket.WebSocket`
+        class is still coupled a bit to the socket interface,
+        we have to override a little more than necessary
+        to play nice with the :pep:`3156` interface. Hopefully,
+        some day this will be cleaned out.
+        """
+        _WebSocket.__init__(self, None)
+        self.started = False
+        self.proto = proto
+
+    @property
+    def local_address(self):
+        """
+        Local endpoint address as a tuple
+        """
+        if not self._local_address:
+            self._local_address = self.proto.reader.transport.get_extra_info('sockname')
+            if len(self._local_address) == 4:
+                self._local_address = self._local_address[:2]
+        return self._local_address
+
+    @property
+    def peer_address(self):
+        """
+        Peer endpoint address as a tuple
+        """
+        if not self._peer_address:
+            self._peer_address = self.proto.reader.transport.get_extra_info('peername')
+            if len(self._peer_address) == 4:
+                self._peer_address = self._peer_address[:2]
+        return self._peer_address
+
+    def once(self):
+        """
+        The base class directly is used in conjunction with
+        the :class:`ws4py.manager.WebSocketManager` which is
+        not actually used with the asyncio implementation
+        of ws4py. So let's make it clear it shan't be used.
+        """
+        raise NotImplemented()
+
+    def close_connection(self):
+        """
+        Close the underlying transport
+        """
+        @asyncio.coroutine
+        def closeit():
+            yield from self.proto.writer.drain()
+            self.proto.writer.close()
+        asyncio.async(closeit())
+
+    def _write(self, data):
+        """
+        Write to the underlying transport
+        """
+        @asyncio.coroutine
+        def sendit(data):
+            self.proto.writer.write(data)
+            yield from self.proto.writer.drain()
+        asyncio.async(sendit(data))
+
+    @asyncio.coroutine
+    def run(self):
+        """
+        Coroutine that runs until the websocket
+        exchange is terminated. It also calls the
+        `opened()` method to indicate the exchange
+        has started.
+        """
+        self.started = True
+        try:
+            self.opened()
+            reader = self.proto.reader
+            while True:
+                data = yield from reader.read(self.reading_buffer_size)
+                if not self.process(data):
+                    return False
+        finally:
+            self.terminate()
+
+        return True
+
+class EchoWebSocket(WebSocket):
+    def received_message(self, message):
+        """
+        Automatically sends back the provided ``message`` to
+        its originating endpoint.
+        """
+        self.send(message.data, message.is_binary)