Source.from_tcp fixes

Running the network word count example showed an attempt to await a list
object in the Tornado framework. I tracked it down to the
EmitServer.handle_stream assuming the result of source._emit was
awaitable. Pivoting on inspect.isawaitable there makes the example as
written work.

I added a try/except block to trap the keyboard interrupt and suppress
the stack spew to the example.

Author: dwskoog

examples/network_wordcount.py

@@ -18,4 +18,9 @@
 )
 
 s.start()
-time.sleep(600)
+
+try:
+    while True:
+        time.sleep(600)
+except KeyboardInterrupt:
+    pass

streamz/sources.py

@@ -3,6 +3,8 @@
 import queue
 import os
 import time
+from inspect import isawaitable
+
 from tornado import gen
 import weakref
 
@@ -252,7 +254,9 @@ async def handle_stream(self, stream, address):
                 while not self.source.stopped:
                     try:
                         data = await stream.read_until(self.source.delimiter)
-                        await self.source._emit(data)
+                        result = self.source._emit(data)
+                        if isawaitable(result):
+                            await result
                     except StreamClosedError:
                         break
 

streamz/tests/test_sources.py

@@ -47,6 +47,37 @@ def test_tcp():
         sock2.close()
 
 
+@flaky(max_runs=3, min_passes=1)
+def test_tcp_word_count_example():
+    port = 9999
+    s = Source.from_tcp(port)
+    out = s.map(bytes.split).flatten().frequencies().sink_to_list()
+    s.start()
+    wait_for(lambda: s.server is not None, 2, period=0.02)
+
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+        sock.connect(("localhost", port))
+        sock.send(b'data\n')
+
+    with (socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock,
+          socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock2):
+        sock.connect(("localhost", port))
+        sock2.connect(("localhost", port))
+        sock.send(b'data\n')
+        # regression test a bug in from_tcp where a second packet from
+        # the same socket is dropped due to the socket handler dying
+        sock.send(b'data\n')
+        sock2.send(b'data2\n')
+
+    expected = [{b"data": 1}, {b"data": 2}, {b"data": 3}, {b"data": 3, b"data2": 1}]
+
+    def fail_func():
+        assert out == expected
+
+    wait_for(lambda: out == expected, 2, fail_func=fail_func, period=0.01)
+
+
+
 @flaky(max_runs=3, min_passes=1)
 @gen_test(timeout=60)
 def test_tcp_async():