Issue Gandi#7: drop timeout, use a non-blocking socket.

Signed-off-by: Arthur Gautier <baloo@gandi.net>

Author: Arthur Gautier

tox.ini

@@ -1,7 +1,6 @@
 [tox]
 envlist =
     {py27,py3}-{lowest,release},
-# No python3 yet, not supported by thrift, see https://issues.apache.org/jira/browse/THRIFT-1857
     pep8
 #    pylint,
 #    coverage

zipkin/client.py

@@ -1,8 +1,14 @@
 import logging
-from socket import timeout
+import socket
+import errno
+import struct
+from io import BytesIO
 
-from thriftpy.transport import TTransportException, TFramedTransportFactory
-from thriftpy.rpc import make_client
+from thriftpy.transport import TTransportException, TFramedTransportFactory,\
+    TSocket
+from thriftpy.protocol import TBinaryProtocolFactory
+from thriftpy.protocol.binary import write_message_begin, write_val
+from thriftpy.thrift import TClient, TType, TMessageType
 
 from .util import base64_thrift_formatter
 from .scribe import scribe_thrift
@@ -12,24 +18,109 @@
 
 CONNECTION_RETRIES = [1, 10, 20, 50, 100, 200, 400, 1000]
 
+try:
+    MSG_NOSIGNAL = socket.MSG_NOSIGNAL
+except:
+    MSG_NOSIGNAL = 16384  # python2
+
+
+class TNonBlockingSocket(TSocket):
+    def _init_sock(self):
+        super(TNonBlockingSocket, self)._init_sock()
+
+        # 1M sendq buffer
+        self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 1024*1024)
+
+        self.sock.setblocking(0)
+
+    def open(self):
+        self._init_sock()
+
+        addr = self.unix_socket or (self.host, self.port)
+        status = self.sock.connect_ex(addr)
+
+        if status not in [errno.EINPROGRESS, errno.EALREADY]:
+            raise IOError("connection attempt on a non-clean socket",
+                          errno.errorcode[status])
+
+    def write(self, buff):
+        # First ensure our incoming end is always empty
+        self.read_all()
+
+        # Then actually try to write to socket
+        try:
+            # We are a library, we can't just set sighandlers. But we don't
+            # want SIGPIPE if peer has gone away either. We better set
+            # MSG_NOSIGNAL to avoid that.
+            # If peer has disconnected, then a errno.EPIPE will raise, and
+            # will be catched on the uppper layer
+
+            self.sock.sendall(buff, MSG_NOSIGNAL)
+        except socket.error as e:
+            if e.errno not in [errno.EINPROGRESS,  # Not connected yet
+                               errno.EWOULDBLOCK]:  # write buffer full
+                # In all other cases, raise.
+                raise
+
+            # If not yet connected or write buffer is full, silently drop.
+
+    def read_all(self):
+        """
+        Flush incoming buffer
+        """
+        try:
+            while True:  # socket.error.errno.EAGAIN will exit this
+                self.sock.recv(1024)
+        except socket.error as e:
+            # if EAGAIN or EWOULDBLOCK, then there is nothing to read
+            if e.errno not in [errno.EAGAIN, errno.EWOULDBLOCK]:
+                # Otherwise that's an error.
+                raise
+
+            return  # No more data to read, or connection is not ready
+
+    def read(self, _):
+        """
+        Mock response, we don't care about results. We never actually read
+        them. But we don't want client to wait for server to reply.
+        """
+        buffer = BytesIO()
+        seq_id = 0  # Sequence id is never compared to message.
+        write_message_begin(buffer, 'Log_result', TMessageType.REPLY, seq_id)
+
+        response = scribe_thrift.Scribe.Log_result(
+            success=scribe_thrift.ResultCode.OK)
+        write_val(buffer, TType.STRUCT, response)
+
+        out = buffer.getvalue()
+
+        # Framed message, starts with length of message.
+        return struct.pack('!i', len(out)) + out
+
+
+def make_client(service, host, port,
+                proto_factory=TBinaryProtocolFactory(),
+                trans_factory=TFramedTransportFactory()):
+
+    socket = TNonBlockingSocket(host, port)
+    transport = trans_factory.get_transport(socket)
+    protocol = proto_factory.get_protocol(transport)
+    transport.open()
+    return TClient(service, protocol)
+
 
 class Client(object):
 
     host = None
     port = 9410
     _client = None
     _connection_attempts = 0
-    # This is used in TSocket which divides by 1000 before passing it to
-    # python socket. This is ms, do not trust socket.settimeout documentation.
-    timeout = 50
 
     @classmethod
     def configure(cls, settings, prefix):
         cls.host = settings.get(prefix + 'collector')
         if prefix + 'collector.port' in settings:
             cls.port = int(settings[prefix + 'collector.port'])
-        if prefix + 'collector.timeout' in settings:
-            cls.timeout = int(settings[prefix + 'collector.timeout'])
 
     @classmethod
     def get_connection(cls):
@@ -47,8 +138,7 @@ def get_connection(cls):
             try:
                 cls._client = make_client(
                     scribe_thrift.Scribe, host=cls.host,
-                    port=cls.port, timeout=cls.timeout,
-                    trans_factory=TFramedTransportFactory())
+                    port=cls.port)
 
                 cls._connection_attempts = 0
             except TTransportException:
@@ -79,10 +169,6 @@ def log(cls, trace):
                 cls._client = None
                 logger.error('EOFError while logging a trace on zipkin '
                              'collector %s:%d' % (cls.host, cls.port))
-            except timeout:
-                cls._client = None
-                logger.error('timeout when sending data or connecting to '
-                             'collector %s:%d' % (cls.host, cls.port))
             except Exception:
                 cls._client = None
                 logger.exception('Unknown Exception while logging a trace on '
@@ -91,6 +177,11 @@ def log(cls, trace):
         else:
             logger.warn("Can't log zipkin trace, not connected")
 
+    @classmethod
+    def disconnect(cls):
+        if cls._client:
+            cls._client.close()
+
 
 def log(trace):
     Client.log(trace)

zipkin/tests/test_client.py

@@ -0,0 +1,80 @@
+import unittest
+import threading
+import socket
+import select
+
+from ..client import Client
+from ..models import Trace, Annotation
+
+
+class Sinkhole(threading.Thread):
+    def __init__(self, *args, **kwargs):
+        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        self.sock.bind(('127.0.0.1', 0))
+        self.port = self.sock.getsockname()[1]
+        self.sock.listen(5)
+
+        self.stopped = False
+
+        super(Sinkhole, self).__init__()
+
+        self.has_received_data = False
+        self.has_data = threading.Lock()
+        self.has_data.acquire()
+
+    def run(self):
+        sockets = [self.sock]
+
+        while not self.stopped:
+            readable, _, _ = select.select(sockets, [], [])
+
+            for s in readable:
+                if s is self.sock:
+                    connection, _ = s.accept()
+                    connection.setblocking(0)
+                    sockets.append(connection)
+                else:
+                    data = s.recv(1024)
+                    if not data:
+                        sockets.remove(s)
+                        s.close()
+                    else:
+                        self.has_received_data = True
+                        self.has_data.release()
+
+        self.sock.close()
+
+    def stop(self):
+        self.stopped = True
+        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        s.connect_ex(('127.0.0.1', self.port))
+        s.close()
+        self.join()
+
+
+class TestClient(unittest.TestCase):
+    def setUp(self):
+        self.server = Sinkhole()
+        self.server.start()
+
+    def tearDown(self):
+        self.server.stop()
+
+    def test_simple(self):
+        Client.configure({
+          'collector': '127.0.0.1',
+          'collector.port': self.server.port,
+        }, '')
+
+        trace = Trace('test')
+        trace.record(Annotation.string('foo', 'bar'))
+        Client.get_connection()
+
+        import time
+        time.sleep(1)  # Just to give it time to connect
+
+        Client.log(trace)
+        self.server.has_data.acquire()
+        self.server.has_data.release()
+        self.assertTrue(self.server.has_received_data)
+        Client.disconnect()