Rework BatchedSend logic

distributed/batched.py

@@ -10,7 +10,7 @@
 from tornado.iostream import StreamClosedError
 from tornado.ioloop import PeriodicCallback, IOLoop
 
-from .core import read, write
+from .core import read, write, flush
 from .utils import ignoring, log_errors
 
 
@@ -112,6 +112,7 @@ def close(self, ignore_closed=False):
             if self.buffer:
                 self.buffer, payload = [], self.buffer
                 yield write(self.stream, payload)
+            yield flush(self.stream)
         except StreamClosedError:
             if not ignore_closed:
                 raise

distributed/core.py

@@ -186,7 +186,7 @@ def handle_stream(self, stream, address):
                     break
         finally:
             try:
-                stream.close()
+                yield close(stream)
             except Exception as e:
                 logger.warn("Failed while closing writer",  exc_info=True)
         logger.info("Close connection from %s:%d to %s", address[0], address[1],
@@ -229,11 +229,36 @@ def write(stream, msg):
     stream.write(b''.join(lengths))
 
     for frame in frames:
+        # Can't wait for the write() Future as it may be lost
+        # ("If write is called again before that Future has resolved,
+        #   the previous future will be orphaned and will never resolve")
         stream.write(frame)
 
     yield gen.moment
 
 
+@gen.coroutine
+def flush(stream):
+    """Flush the stream's output buffer.
+    This is recommended before closing the stream.
+    """
+    if stream.writing():
+        yield stream.write(b'')
+
+
+@gen.coroutine
+def close(stream):
+    """Close a stream safely.
+    """
+    if not stream.closed():
+        try:
+            flush(stream)
+        except StreamClosedError:
+            pass
+        finally:
+            stream.close()
+
+
 def pingpong(stream):
     return b'pong'
 
@@ -295,7 +320,7 @@ def send_recv(stream=None, arg=None, ip=None, port=None, addr=None, reply=True,
     else:
         response = None
     if kwargs.get('close'):
-        stream.close()
+        close(stream)
     raise gen.Return(response)
 
 
@@ -398,7 +423,7 @@ def close_streams(self):
         for stream in self.streams:
             if stream and not stream.closed():
                 try:
-                    stream.close()
+                    close(stream)
                 except (OSError, IOError, StreamClosedError):
                     pass
         self.streams.clear()
@@ -551,15 +576,15 @@ def collect(self):
                     self.open, self.active)
         for streams in list(self.available.values()):
             for stream in streams:
-                stream.close()
+                close(stream)
 
     def close(self):
         for streams in list(self.available.values()):
             for stream in streams:
-                stream.close()
+                close(stream)
         for streams in list(self.occupied.values()):
             for stream in streams:
-                stream.close()
+                close(stream)
 
 
 def coerce_to_address(o, out=str):

distributed/scheduler.py

@@ -29,7 +29,7 @@
 
 from .batched import BatchedSend
 from .config import config
-from .core import (rpc, connect, read, write, MAX_BUFFER_SIZE,
+from .core import (rpc, connect, read, write, close, MAX_BUFFER_SIZE,
         Server, send_recv, coerce_to_address, error_message)
 from .utils import (All, ignoring, clear_queue, get_ip, ignore_exceptions,
         ensure_ip, get_fileno_limit, log_errors, key_split, mean,
@@ -410,7 +410,7 @@ def finished(self):
     def close_streams(self):
         """ Close all active IOStreams """
         for stream in self.streams.values():
-            stream.stream.close()
+            close(stream.stream)
         self.rpc.close()
 
     @gen.coroutine
@@ -711,8 +711,7 @@ def remove_worker(self, stream=None, address=None, safe=False):
                 return 'already-removed'
             with ignoring(AttributeError):
                 stream = self.worker_streams[address].stream
-                if not stream.closed():
-                    stream.close()
+                close(stream)
 
             host, port = address.split(':')
 
@@ -1141,8 +1140,7 @@ def worker_stream(self, worker):
                 import pdb; pdb.set_trace()
             raise
         finally:
-            if not stream.closed():
-                stream.close()
+            close(stream)
             self.remove_worker(address=worker)
 
     def correct_time_delay(self, worker, msg):

distributed/tests/test_nanny.py

@@ -12,7 +12,7 @@
 from tornado import gen
 
 from distributed import Nanny, rpc, Scheduler
-from distributed.core import connect, read, write
+from distributed.core import connect, read, write, close
 from distributed.protocol.pickle import dumps, loads
 from distributed.utils import ignoring
 from distributed.utils_test import gen_cluster
@@ -120,7 +120,7 @@ def test_monitor_resources(s):
         assert isinstance(msg, dict)
         assert {'cpu_percent', 'memory_percent'}.issubset(msg)
 
-    stream.close()
+    close(stream)
     yield n._close()
     s.stop()
 

distributed/tests/test_scheduler.py

@@ -20,7 +20,7 @@
 import pytest
 
 from distributed import Nanny, Worker
-from distributed.core import connect, read, write, rpc
+from distributed.core import connect, read, write, close, rpc
 from distributed.scheduler import (validate_state, decide_worker,
         Scheduler)
 from distributed.client import _wait
@@ -402,7 +402,7 @@ def test_scheduler(s, a, b):
             break
 
     write(stream, {'op': 'close'})
-    stream.close()
+    close(stream)
 
 
 @gen_cluster()
@@ -467,7 +467,7 @@ def test_server(s, a, b):
     assert msg == {'op': 'stream-closed'}
     with pytest.raises(StreamClosedError):
         yield readone(stream)
-    stream.close()
+    close(stream)
 
 
 @gen_cluster()
@@ -547,7 +547,7 @@ def func(scheduler):
         expected = s.processing, s.stacks
         assert cloudpickle.loads(response) == expected
 
-    stream.close()
+    close(stream)
 
 
 @gen_cluster()
@@ -573,7 +573,7 @@ def teardown(scheduler, state):
         response = yield read(stream)
         assert response == 'OK'
 
-    stream.close()
+    close(stream)
     start = time()
     while not hasattr(s, 'flag'):
         yield gen.sleep(0.01)
@@ -599,7 +599,7 @@ def func(scheduler):
         response = yield read(stream)
         assert response == True
 
-    stream.close()
+    close(stream)
 
 
 @gen_test(timeout=None)

distributed/utils.py

@@ -63,7 +63,7 @@ def get_ip(host='8.8.8.8', port=80):
 def ignoring(*exceptions):
     try:
         yield
-    except exceptions:
+    except exceptions as e:
         pass
 
 

distributed/utils_test.py

@@ -20,7 +20,7 @@
 from tornado.ioloop import IOLoop, TimeoutError
 from tornado.iostream import StreamClosedError
 
-from .core import connect, read, write, rpc
+from .core import connect, read, write, close, rpc
 from .utils import ignoring, log_errors, sync
 import pytest
 
@@ -303,7 +303,7 @@ def disconnect(ip, port):
         yield write(stream, {'op': 'terminate', 'close': True})
         response = yield read(stream)
     finally:
-        stream.close()
+        yield close(stream)
 
 
 import pytest