main.py

from kazoo.client import KazooClient
import argparse
import subprocess
from functools import partial
from kazoo.exceptions import NoNodeException

parser = argparse.ArgumentParser(description='Lab 7 - Zookeeper')
parser.add_argument("--port", type=int, default=2181, help="Zookeeper port")
parser.add_argument("--exec", type=str, default="python gui.py", help="Command to start other process")
args = parser.parse_args()

WATCHED_NODE = "/a"

zk = KazooClient(hosts=f"127.0.0.1:{args.port}")
process = None

children_set = set()


# avoid first call on ChildrenWatch when children are registered and avoid displaying nodes count when creating /a node
def avoid_first_function_call(function):
    first_call = True

    def closure(*args_, **kwargs):
        nonlocal first_call

        if first_call:
            first_call = False
            return

        return function(*args_, **kwargs)

    return closure


@zk.DataWatch(WATCHED_NODE)
def watch_node(_data, _stat, event):
    if event is None:
        return

    global process

    if event.type == "CREATED":
        process = subprocess.Popen(args.exec.split(" "))
        zk.ChildrenWatch(WATCHED_NODE)(partial(avoid_first_function_call(watch_children), path=WATCHED_NODE))
    elif event.type == "DELETED":
        if process is not None:
            process.kill()
            process = None


def watch_children(children, path):
    for child in children:
        child_path = f"{path}/{child}"
        if child_path not in children_set:
            zk.ChildrenWatch(child_path)(partial(avoid_first_function_call(watch_children), path=child_path))
            children_set.add(child_path)

    print(f"Nodes count: {get_nodes_count(WATCHED_NODE)}")


def get_nodes_count(node):
    try:
        children = zk.get_children(node)

        count = 1
        for child in children:
            count += get_nodes_count(f"{node}/{child}")

        return count - 1
    except NoNodeException:
        return 0


def print_tree(node, depth=0):
    try:
        children = zk.get_children(node)
        print("  " * depth + "/" + node.strip("/").split("/")[-1])
        for child in children:
            print_tree(f"{node}/{child}", depth + 1)
    except NoNodeException:
        return


if __name__ == '__main__':
    zk.start()

    while True:
        command = input()
        if command == "quit":
            break
        elif command == "print":
            print_tree(WATCHED_NODE)

    zk.stop()
    zk.close()