Address review comments

distributed/batched.py

@@ -24,6 +24,9 @@ class BatchedSend(object):
     more than one message every interval milliseconds.  We send lists of
     messages.
 
+    Batching several messages at once helps performance when sending
+    a myriad of tiny messages.
+
     Example
     -------
     >>> stream = yield connect(ip, port)
@@ -74,12 +77,12 @@ def _background_send(self):
                 continue
             payload, self.buffer = self.buffer, []
             self.batch_count += 1
+            self.next_deadline = self.loop.time() + self.interval
             try:
                 yield write(self.stream, payload)
             except Exception:
                 logger.exception("Error in batched write")
                 break
-            self.next_deadline = self.loop.time() + self.interval
 
         self.stopped.set()
 
@@ -93,7 +96,9 @@ def send(self, msg):
 
         self.message_count += 1
         self.buffer.append(msg)
-        self.waker.set()
+        # Avoid spurious wakeups if possible
+        if self.next_deadline is None:
+            self.waker.set()
 
     @gen.coroutine
     def close(self, ignore_closed=False):
@@ -105,10 +110,6 @@ def close(self, ignore_closed=False):
         yield self.stopped.wait()
         try:
             if self.buffer:
-                if self.next_deadline is not None:
-                    delay = self.next_deadline - self.loop.time()
-                    if delay > 0:
-                        yield gen.sleep(delay)
                 self.buffer, payload = [], self.buffer
                 yield write(self.stream, payload)
         except StreamClosedError:

distributed/core.py

@@ -231,23 +231,14 @@ def write(stream, msg):
             logger.exception(e)
             raise
 
-        futures = []
-
         lengths = ([struct.pack('Q', len(frames))] +
                    [struct.pack('Q', len(frame)) for frame in frames])
-        futures.append(stream.write(b''.join(lengths)))
+        stream.write(b''.join(lengths))
 
         for frame in frames[:-1]:
-            futures.append(stream.write(frame))
-
-        futures.append(stream.write(frames[-1]))
+            stream.write(frame)
 
-        while stream._write_buffer:
-            try:
-                yield gen.with_timeout(timedelta(seconds=0.01), futures[-1])
-                break
-            except gen.TimeoutError:
-                pass
+        yield stream.write(frames[-1])
 
 
 def pingpong(stream):