kafka fixes

Author: Martin Durant

streamz/core.py

@@ -512,8 +512,10 @@ def disconnect(self, downstream):
 
     @property
     def upstream(self):
-        if len(self.upstreams) != 1:
+        if len(self.upstreams) > 1:
             raise ValueError("Stream has multiple upstreams")
+        elif len(self.upstreams) == 0:
+            return None
         else:
             return self.upstreams[0]
 
@@ -535,6 +537,13 @@ def remove(self, predicate):
         """ Only pass through elements for which the predicate returns False """
         return self.filter(lambda x: not predicate(x))
 
+    def stop(self):
+        """Call on any stream node to halt all upstream sources"""
+        prev, s = self.upstream, self
+        while s:
+            prev, s = s, s.upstream
+        prev.stopped = True
+
     @property
     def scan(self):
         return self.accumulate

streamz/sources.py

@@ -3,6 +3,7 @@
 import os
 import time
 from tornado import gen
+import weakref
 
 from .core import Stream, convert_interval, RefCounter
 
@@ -446,6 +447,7 @@ def start(self):
             self.stopped = False
             self.consumer = ck.Consumer(self.cpars)
             self.consumer.subscribe(self.topics)
+            weakref.finalize(self, self.consumer.close)
             tp = ck.TopicPartition(self.topics[0], 0, 0)
 
             # blocks for consumer thread to come up
@@ -496,7 +498,7 @@ def commit(_part):
 
         @gen.coroutine
         def checkpoint_emit(_part):
-            ref = RefCounter(cb=lambda: commit(_part))
+            ref = RefCounter(cb=lambda: commit(_part), loop=self.loop)
             yield self._emit(_part, metadata=[{'ref': ref}])
 
         if self.npartitions is None:
@@ -521,61 +523,58 @@ def checkpoint_emit(_part):
                     self.positions[tp.partition] = tp.offset
                 break
 
-        try:
-            while not self.stopped:
-                out = []
-
-                if self.refresh_partitions:
-                    kafka_cluster_metadata = self.consumer.list_topics(self.topic)
-                    if self.engine == "cudf":  # pragma: no cover
-                        new_partitions = len(kafka_cluster_metadata[self.topic.encode('utf-8')])
-                    else:
-                        new_partitions = len(kafka_cluster_metadata.topics[self.topic].partitions)
-                    if new_partitions > self.npartitions:
-                        self.positions.extend([-1001] * (new_partitions - self.npartitions))
-                        self.npartitions = new_partitions
+        while not self.stopped:
+            out = []
 
-                for partition in range(self.npartitions):
-                    tp = ck.TopicPartition(self.topic, partition, 0)
-                    try:
-                        low, high = self.consumer.get_watermark_offsets(
-                            tp, timeout=0.1)
-                    except (RuntimeError, ck.KafkaException):
-                        continue
-                    self.started = True
-                    if 'auto.offset.reset' in self.consumer_params.keys():
-                        if self.consumer_params['auto.offset.reset'] == 'latest' and \
-                                self.positions[partition] == -1001:
-                            self.positions[partition] = high
-                    current_position = self.positions[partition]
-                    lowest = max(current_position, low)
-                    if high > lowest + self.max_batch_size:
-                        high = lowest + self.max_batch_size
-                    if high > lowest:
-                        out.append((self.consumer_params, self.topic, partition,
-                                    self.keys, lowest, high - 1))
+            if self.refresh_partitions:
+                kafka_cluster_metadata = self.consumer.list_topics(self.topic)
+                if self.engine == "cudf":  # pragma: no cover
+                    new_partitions = len(kafka_cluster_metadata[self.topic.encode('utf-8')])
+                else:
+                    new_partitions = len(kafka_cluster_metadata.topics[self.topic].partitions)
+                if new_partitions > self.npartitions:
+                    self.positions.extend([-1001] * (new_partitions - self.npartitions))
+                    self.npartitions = new_partitions
+
+            for partition in range(self.npartitions):
+                tp = ck.TopicPartition(self.topic, partition, 0)
+                try:
+                    low, high = self.consumer.get_watermark_offsets(
+                        tp, timeout=0.1)
+                except (RuntimeError, ck.KafkaException):
+                    continue
+                self.started = True
+                if 'auto.offset.reset' in self.consumer_params.keys():
+                    if self.consumer_params['auto.offset.reset'] == 'latest' and \
+                            self.positions[partition] == -1001:
                         self.positions[partition] = high
-                self.consumer_params['auto.offset.reset'] = 'earliest'
+                current_position = self.positions[partition]
+                lowest = max(current_position, low)
+                if high > lowest + self.max_batch_size:
+                    high = lowest + self.max_batch_size
+                if high > lowest:
+                    out.append((self.consumer_params, self.topic, partition,
+                                self.keys, lowest, high - 1))
+                    self.positions[partition] = high
+            self.consumer_params['auto.offset.reset'] = 'earliest'
+
+            for part in out:
+                yield self.loop.add_callback(checkpoint_emit, part)
 
-                for part in out:
-                    yield self.loop.add_callback(checkpoint_emit, part)
-
-                else:
-                    yield gen.sleep(self.poll_interval)
-        finally:
-            self.consumer.unsubscribe()
-            self.consumer.close()
+            else:
+                yield gen.sleep(self.poll_interval)
 
     def start(self):
         import confluent_kafka as ck
         if self.engine == "cudf":  # pragma: no cover
             from custreamz import kafka
 
         if self.stopped:
-            if self.engine == "cudf": # pragma: no cover
+            if self.engine == "cudf":  # pragma: no cover
                 self.consumer = kafka.Consumer(self.consumer_params)
             else:
                 self.consumer = ck.Consumer(self.consumer_params)
+            weakref.finalize(self, self.consumer.close)
             self.stopped = False
             tp = ck.TopicPartition(self.topic, 0, 0)
 

streamz/tests/test_kafka.py

@@ -116,9 +116,9 @@ def test_from_kafka():
         stream = Stream.from_kafka([TOPIC], ARGS, asynchronous=True)
         out = stream.sink_to_list()
         stream.start()
-        yield gen.sleep(0.1)  # for loop to run
+        yield gen.sleep(1.1)  # for loop to run
         for i in range(10):
-            yield gen.sleep(0.2)
+            yield gen.sleep(0.1)  # small pause ensures correct ordering
             kafka.produce(TOPIC, b'value-%d' % i)
         kafka.flush()
         # it takes some time for messages to come back out of kafka
@@ -168,7 +168,9 @@ def test_from_kafka_thread():
         stream = Stream.from_kafka([TOPIC], ARGS)
         out = stream.sink_to_list()
         stream.start()
+        yield gen.sleep(1.1)
         for i in range(10):
+            yield gen.sleep(0.1)
             kafka.produce(TOPIC, b'value-%d' % i)
         kafka.flush()
         # it takes some time for messages to come back out of kafka
@@ -231,8 +233,10 @@ def test_kafka_dask_batch(c, s, w1, w2):
             kafka.produce(TOPIC, b'value-%d' % i)
         kafka.flush()
         yield await_for(lambda: any(out), 10, period=0.2)
-        assert {'key':None, 'value':b'value-1'} in out[0]
-        stream.upstream.stopped = True
+        assert {'key': None, 'value': b'value-1'} in out[0]
+        stream.stop()
+        yield gen.sleep(0)
+        stream.upstream.upstream.consumer.close()
 
 
 def test_kafka_batch_npartitions():
@@ -551,11 +555,12 @@ def test_kafka_batch_checkpointing_async_nodes_2():
         assert committed3[1].offset == 1
 
 
+@flaky(max_runs=3, min_passes=1)
 def test_kafka_checkpointing_auto_offset_reset_latest():
-    '''
+    """
     Testing whether checkpointing works as expected with multiple topic partitions and
     with auto.offset.reset configuration set to latest (also default).
-    '''
+    """
     j = random.randint(0, 10000)
     ARGS = {'bootstrap.servers': 'localhost:9092',
             'group.id': 'streamz-test%i' % j,