Multi-kafka for high availability

Start by running the provided asyncio example with two from_kafka_batched sources and an a.union(b) sink, then stop one Kafka cluster to reproduce the freeze. Trace the asynchronous Kafka source and union behavior to identify where progress stops. Done means the combined stream continues handling the available cluster when the other cluster is stopped, with coverage for the failure case.

We are setting up Kafka in multiple data centers to achieve high availability. I expected that Streamz could do a union between both of these clusters. Everything works as expected, but when I stop one of the Kafka clusters, the stream freezes. I'm thinking there is a callback somewhere that is blocking operations?

import asyncio

import confluent_kafka as ck
from streamz import from_kafka_batched


async def produce(broker, cluster_id, topic):
    p = ck.Producer({'bootstrap.servers': broker})

    for i in range(100000):
        p.poll(0)
        p.produce(topic, 'cluster-{} value-{}'.format(cluster_id, i))
        p.flush(timeout=0.5)
        await asyncio.sleep(1)


async def main():
    topic = 'test-topic'
    args1 = {'bootstrap.servers': 'localhost:9092', 'group.id': 'test-group'}
    args2 = {'bootstrap.servers': 'localhost:9093', 'group.id': 'test-group'}
    a = from_kafka_batched(topic, args1, npartitions=1, start=True,
                           asynchronous=True)
    b = from_kafka_batched(topic, args2, npartitions=1, start=True,
                           asynchronous=True)

    a.union(b).sink(print)

    asyncio.create_task(produce('localhost:9092', 1, topic))
    asyncio.create_task(produce('localhost:9093', 2, topic))

    while True:
        await asyncio.sleep(1)


if __name__ == '__main__':
    asyncio.run(main())