Message size check update

arachnado/rpc/data.py

@@ -30,12 +30,7 @@ class DataRpcWebsocketHandler(RpcWebsocketHandler):
     max_msg_size = 2**20
 
     def _send_event(self, data):
-        message = json_encode(data)
-        # if message size is higher then ws connection can be dropped without proper message
-        if sys.getsizeof(message) < self.max_msg_size or not self.max_msg_size:
-            return super(DataRpcWebsocketHandler, self).write_event(data)
-        else:
-            logger.info("Message size exceeded. Message wasn't sent.")
+        return super(DataRpcWebsocketHandler, self).write_event(data, max_message_size=self.max_msg_size)
 
     def init_heartbeat(self, update_delay):
         if update_delay > 0 and not self.heartbeat_data:

arachnado/rpc/ws.py

@@ -1,3 +1,4 @@
+import sys
 import json
 import logging
 import six
@@ -30,13 +31,16 @@ def send_data(self, data):
         self.write_event(data)
 
     @gen.coroutine
-    def write_event(self, data):
+    def write_event(self, data, max_message_size=0):
         if isinstance(data, six.string_types):
             message = data
         else:
             message = json_encode(data)
         try:
-            self.write_message(message)
+            if sys.getsizeof(message) < max_message_size or not max_message_size:
+                self.write_message(message)
+            else:
+                logger.info("Message size exceeded. Message wasn't sent.")
         except websocket.WebSocketClosedError:
             pass