keep_alive.py

import time
from threading import Thread

from kazoo.client import KazooState


def create_nodes_after_reconnect(zk):
    time.sleep(0.1)
    for (service_path, data) in zk.ka_nodes.items():
        register_service(zk, service_path, data, keep_alive=True)


def start_ka_handler(zk):
    zk.ka_nodes = {}

    def connection_handler(zk_state):
        if getattr(zk, 'ka_handler_started', False):
            if zk_state == KazooState.CONNECTED:
                Thread(target=create_nodes_after_reconnect, args=(zk,)).start()
        else:
            setattr(zk, 'ka_handler_started', True)

    zk.add_listener(connection_handler)