Add proper handling of the send, tracking if data was actually sent in fluent#77

arcivanov · arcivanov · commit 161d76e2f791 · 2017-12-04T14:28:56.000-05:00
Use TCP_NODELAY to detect disconnect early Fix connection shutdown cleanup in fluent#103 Make locking less verbose via with idiom fixes fluent#77, fluent#103
diff --git a/fluent/asynchandler.py b/fluent/asynchandler.py
@@ -2,7 +2,6 @@
 
 from fluent import asyncsender
 from fluent import handler
-from fluent.handler import FluentRecordFormatter
 
 
 class FluentHandler(handler.FluentHandler):
@@ -14,5 +13,7 @@ def getSenderClass(self):
         return asyncsender.FluentSender
 
     def close(self):
-        self.sender.close()
-        super(FluentHandler, self).close()
+        try:
+            self.sender.close()
+        finally:
+            super(FluentHandler, self).close()
diff --git a/fluent/asyncsender.py b/fluent/asyncsender.py
@@ -1,8 +1,10 @@
 # -*- coding: utf-8 -*-
 
 from __future__ import print_function
+
 import threading
 import time
+
 try:
     from queue import Queue, Full, Empty
 except ImportError:
@@ -11,28 +13,30 @@
 from fluent import sender
 from fluent.sender import EventTime
 
+__all__ = ["EventTime", "FluentSender"]
+
 _global_sender = None
 
 DEFAULT_QUEUE_TIMEOUT = 0.05
 
 
-def _set_global_sender(sender):
+def _set_global_sender(sender):  # pragma: no cover
     """ [For testing] Function to set global sender directly
     """
     global _global_sender
     _global_sender = sender
 
 
-def setup(tag, **kwargs):
+def setup(tag, **kwargs):  # pragma: no cover
     global _global_sender
     _global_sender = FluentSender(tag, **kwargs)
 
 
-def get_global_sender():
+def get_global_sender():  # pragma: no cover
     return _global_sender
 
 
-def close():
+def close():  # pragma: no cover
     get_global_sender().close()
 
 
@@ -69,9 +73,8 @@ def run(self):
                 bytes_ = self._queue.get(block=True, timeout=self._queue_timeout)
             except Empty:
                 continue
-            self._conn_close_lock.acquire()
-            self._sender._send(bytes_)
-            self._conn_close_lock.release()
+            with self._conn_close_lock:
+                self._sender._send(bytes_)
 
     def close(self, flush=True, discard=True):
         if discard:
@@ -86,14 +89,8 @@ def close(self, flush=True, discard=True):
         self._sender.close()
 
     def _close(self):
-        self._conn_close_lock.acquire()
-        # self._sender.lock.acquire()
-        try:
+        with self._conn_close_lock:
             self._sender._close()
-        finally:
-            # self._sender.lock.release()
-            self._conn_close_lock.release()
-            pass
 
     @property
     def last_error(self):
@@ -103,7 +100,7 @@ def last_error(self):
     def last_error(self, err):
         self._sender.last_error = err
 
-    def clear_last_error(self, _thread_id = None):
+    def clear_last_error(self, _thread_id=None):
         self._sender.clear_last_error(_thread_id=_thread_id)
 
     @property
@@ -133,14 +130,15 @@ def __init__(self,
                  nanosecond_precision=False,
                  msgpack_kwargs=None,
                  queue_timeout=DEFAULT_QUEUE_TIMEOUT,
-                 **kwargs): # This kwargs argument is not used in __init__. This will be removed in the next major version.
+                 **kwargs):  # This kwargs argument is not used in __init__. This will be removed in the next major version.
         super(FluentSender, self).__init__(tag=tag, host=host, port=port, bufmax=bufmax, timeout=timeout,
                                            verbose=verbose, buffer_overflow_handler=buffer_overflow_handler,
                                            nanosecond_precision=nanosecond_precision, msgpack_kwargs=msgpack_kwargs,
                                            **kwargs)
         self._communicator = CommunicatorThread(tag=tag, host=host, port=port, bufmax=bufmax, timeout=timeout,
                                                 verbose=verbose, buffer_overflow_handler=buffer_overflow_handler,
-                                                nanosecond_precision=nanosecond_precision, msgpack_kwargs=msgpack_kwargs,
+                                                nanosecond_precision=nanosecond_precision,
+                                                msgpack_kwargs=msgpack_kwargs,
                                                 queue_timeout=queue_timeout)
         self._communicator.start()
 
@@ -175,7 +173,7 @@ def last_error(self):
     def last_error(self, err):
         self._communicator.last_error = err
 
-    def clear_last_error(self, _thread_id = None):
+    def clear_last_error(self, _thread_id=None):
         self._communicator.clear_last_error(_thread_id=_thread_id)
 
     @property
diff --git a/fluent/sender.py b/fluent/sender.py
@@ -1,34 +1,35 @@
 # -*- coding: utf-8 -*-
 
 from __future__ import print_function
-import struct
+
 import socket
+import struct
 import threading
 import time
 import traceback
 
 import msgpack
 
-
 _global_sender = None
 
 
-def _set_global_sender(sender):
+def _set_global_sender(sender):  # pragma: no cover
     """ [For testing] Function to set global sender directly
     """
     global _global_sender
     _global_sender = sender
 
 
-def setup(tag, **kwargs):
+def setup(tag, **kwargs):  # pragma: no cover
     global _global_sender
     _global_sender = FluentSender(tag, **kwargs)
 
 
-def get_global_sender():
+def get_global_sender():  # pragma: no cover
     return _global_sender
 
-def close():
+
+def close():  # pragma: no cover
     get_global_sender().close()
 
 
@@ -54,7 +55,7 @@ def __init__(self,
                  buffer_overflow_handler=None,
                  nanosecond_precision=False,
                  msgpack_kwargs=None,
-                 **kwargs): # This kwargs argument is not used in __init__. This will be removed in the next major version.
+                 **kwargs):  # This kwargs argument is not used in __init__. This will be removed in the next major version.
 
         self.tag = tag
         self.host = host
@@ -98,8 +99,7 @@ def emit_with_time(self, label, timestamp, data):
         return self._send(bytes_)
 
     def close(self):
-        self.lock.acquire()
-        try:
+        with self.lock:
             if self.pendings:
                 try:
                     self._send_data(self.pendings)
@@ -108,8 +108,6 @@ def close(self):
 
             self._close()
             self.pendings = None
-        finally:
-            self.lock.release()
 
     def _make_packet(self, label, timestamp, data):
         if label:
@@ -122,11 +120,8 @@ def _make_packet(self, label, timestamp, data):
         return msgpack.packb(packet, **self.msgpack_kwargs)
 
     def _send(self, bytes_):
-        self.lock.acquire()
-        try:
+        with self.lock:
             return self._send_internal(bytes_)
-        finally:
-            self.lock.release()
 
     def _send_internal(self, bytes_):
         # buffering
@@ -142,7 +137,6 @@ def _send_internal(self, bytes_):
 
             return True
         except socket.error as e:
-        #except Exception as e:
             self.last_error = e
 
             # close socket
@@ -161,7 +155,13 @@ def _send_data(self, bytes_):
         # reconnect if possible
         self._reconnect()
         # send message
-        self.socket.sendall(bytes_)
+        bytes_to_send = len(bytes_)
+        bytes_sent = 0
+        while bytes_sent < bytes_to_send:
+            sent = self.socket.send(bytes_[bytes_sent:])
+            if sent == 0:
+                raise BrokenPipeError(32, 'broken pipe')
+            bytes_sent += sent
 
     def _reconnect(self):
         if not self.socket:
@@ -172,6 +172,8 @@ def _reconnect(self):
             else:
                 sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                 sock.settimeout(self.timeout)
+                # This might be controversial and may need to be removed
+                sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
                 sock.connect((self.host, self.port))
             self.socket = sock
 
@@ -189,17 +191,22 @@ def last_error(self):
 
     @last_error.setter
     def last_error(self, err):
-        self._last_error_threadlocal.exception  =  err
+        self._last_error_threadlocal.exception = err
 
-    def clear_last_error(self, _thread_id = None):
+    def clear_last_error(self, _thread_id=None):
         if hasattr(self._last_error_threadlocal, 'exception'):
             delattr(self._last_error_threadlocal, 'exception')
 
     def _close(self):
-        if self.socket:
-            self.socket.shutdown(socket.SHUT_RDWR)
-            self.socket.close()
-        self.socket = None
+        try:
+            sock = self.socket
+            if sock:
+                try:
+                    sock.shutdown(socket.SHUT_RDWR)
+                finally:
+                    sock.close()
+        finally:
+            self.socket = None
 
     def __enter__(self):
         return self
diff --git a/tests/test_sender.py b/tests/test_sender.py
@@ -1,8 +1,10 @@
 # -*- coding: utf-8 -*-
 
 from __future__ import print_function
-import unittest
+
 import socket
+import unittest
+
 import msgpack
 
 import fluent.sender
@@ -127,8 +129,9 @@ def test_clear_last_error(self):
 
         self.assertEqual(self._sender.last_error, None)
 
-    @unittest.skip("This test failed with 'TypeError: catching classes that do not inherit from BaseException is not allowed' so skipped")
-    #@patch('fluent.sender.socket')
+    @unittest.skip(
+        "This test failed with 'TypeError: catching classes that do not inherit from BaseException is not allowed' so skipped")
+    # @patch('fluent.sender.socket')
     def test_connect_exception_during_sender_init(self, mock_socket):
         # Make the socket.socket().connect() call raise a custom exception
         mock_connect = mock_socket.socket.return_value.connect
