From 862598f597a1fcbc10bf350d18782929c9d834cc Mon Sep 17 00:00:00 2001 From: Blallo Date: Tue, 5 Feb 2019 15:11:31 +0100 Subject: [PATCH] Improve cli and fix daemon main loop. --- bot_z/cli.py | 47 ++++++++++++++++++++++++++------- bot_z/zdaemon.py | 68 +++++++++++++++++++++++++++++++++++++++--------- 2 files changed, 92 insertions(+), 23 deletions(-) diff --git a/bot_z/cli.py b/bot_z/cli.py index 5a68edc..09de09f 100644 --- a/bot_z/cli.py +++ b/bot_z/cli.py @@ -35,6 +35,9 @@ def _check_name(lifo_path: str, name: T.Optional[str]) -> str: @click.group() @click.option("-d", "--debug", is_flag=True, default=False, help="Enable debug mode.") +@click.option( + "--headless", is_flag=True, default=True, help="Start the clients in headless mode." +) @click.option( "-v", "--verbose", @@ -45,20 +48,25 @@ def _check_name(lifo_path: str, name: T.Optional[str]) -> str: @click.option( "-f", "--fifo", - type=click.Path(), - default="/tmp/bot_z.cmd", + type=click.STRING, + default="bot_z.cmd", help="Path to the control fifo.", ) @click.option( "-w", "--workdir", type=click.Path(exists=True, readable=True, writable=True, resolve_path=True), - default="/tmp/", + default="/tmp", help="The working dir where to launch the daemon and where the lockfile is put.", ) @click.pass_context def cli( - ctx: click.Context, debug: bool, verbose: bool, fifo: str, workdir: str + ctx: click.Context, + debug: bool, + headless: bool, + verbose: bool, + fifo: str, + workdir: str, ) -> None: """ Group cli. @@ -67,8 +75,9 @@ def cli( logger.setLevel(logging.DEBUG) ctx.ensure_object(dict) ctx.obj["debug"] = debug + ctx.obj["headless"] = headless ctx.obj["verbose"] = verbose - ctx.obj["fifo"] = fifo + ctx.obj["fifo"] = os.path.join(workdir, fifo) ctx.obj["workdir"] = workdir ctx.obj["lifo"] = os.path.join(workdir, "botz_open_daemons.list") @@ -124,11 +133,12 @@ def start_daemon_command( base_uri=baseuri, timeout=timeout, proxy=proxy_tuple, - headless=not ctx.obj["debug"], + headless=ctx.obj["headless"], debug=ctx.obj["debug"], foreground=foreground, ) - logger.info("Daemon started.") + if not foreground: + logger.info("Daemon started.") @cli.command("list") @@ -174,6 +184,19 @@ def start_command(ctx: click.Context, name: str) -> None: logger.info("Start sent.") +@cli.command("stop-daemon") +@click.pass_context +def stop_daemon_command(ctx: click.Context) -> None: + """ + Writes on the fifo. Invokes the stop of all the clients and the + subsequent shutdown of the daemon. + """ + logger.debug("Sending the stop-daemon command down the pipe: %s", ctx.obj["fifo"]) + with open(ctx.obj["fifo"], "w") as fifo: + fifo.write(cmd_marshal(name="", cmd="stop-daemon")) + logger.info("Stop-daemon sent.") + + @cli.command("stop") @click.option("-n", "--name", default=None, help="The instance to interact with.") @click.pass_context @@ -186,7 +209,6 @@ def stop_command(ctx: click.Context, name: T.Optional[str]) -> None: logging.info("Stopping instance: %s", name) with open(ctx.obj["fifo"], "w") as fifo: fifo.write(cmd_marshal(name=name, cmd="stop")) - fifo.flush() logger.info("Stop sent.") @@ -201,7 +223,6 @@ def reload_command(ctx: click.Context, name: T.Optional[str]) -> None: name = _check_name(ctx.obj["lifo"], name) with open(ctx.obj["fifo"], "w") as fifo: fifo.write(cmd_marshal(name=name, cmd="reload")) - fifo.flush() logger.info("Reload sent.") @@ -212,7 +233,13 @@ def status_command(ctx: click.Context, name: T.Optional[str]) -> None: """ Writes on the fifo. Queries the status. """ - name = _check_name(ctx.obj["lifo"], name) + try: + name = _check_name(ctx.obj["lifo"], name) + except IndexError: + if len(PLifo.All(ctx.obj["lifo"])) != 0: + raise + logger.warning("No clients registered.") + return _status_command(ctx, name) diff --git a/bot_z/zdaemon.py b/bot_z/zdaemon.py index fed4212..a5e7cd8 100644 --- a/bot_z/zdaemon.py +++ b/bot_z/zdaemon.py @@ -36,6 +36,12 @@ logger.addHandler(syslog) logger.debug("Init at debug") +class StopDaemon(Exception): + """Auxiliary exception to help stop the program in daemon mode.""" + + pass + + def start_daemon( name: str, base_uri: str, @@ -135,7 +141,12 @@ def daemon_process( """ The daemon function. """ + if debug: + logger.setLevel(logging.DEBUG) lifo_path = os.path.join(working_dir, "botz_open_daemons.list") + if os.path.exists(lifo_path): + logger.warning("Lifo (%s) exists. Removing!", lifo_path) + os.remove(lifo_path) starter = functools.partial( start_daemon, @@ -187,8 +198,11 @@ def daemon_process( return with context: logger.debug("Started in background") - listen_commands(lifo_path, fifo_path, cmd_map) - os.remove(lifo_path) + try: + listen_commands(lifo_path, fifo_path, cmd_map) + except StopDaemon: + os.remove(lifo_path) + return def listen_client( @@ -234,24 +248,52 @@ def listen_commands( for cmd_str in fifo: logger.debug("Command received in main loop: %s", cmd_str) cmd = cmd_unmarshal(cmd_str) - if cmd["name"] not in client_chans: - client_chans[cmd["name"]] = queue.Queue(1) - if cmd["name"] in PLifo.All(lifopath): - logger.warning("Name %s is yet used. Not proceeding.", name) - continue - PLifo.Push(lifopath, cmd["name"]) - logger.debug("Client %s has been newly created.", cmd["name"]) - logger.debug("Opening client: %s", cmd["name"]) - q = client_chans[cmd["name"]] + logger.debug("Cmd unmarshalled: %s", cmd) if cmd["cmd"] == "start": + if cmd["name"] not in client_chans: + client_chans[cmd["name"]] = queue.Queue(1) + logger.debug('Queue created for "%s".', cmd["name"]) + if cmd["name"] in PLifo.All(lifopath): + logger.warning("Name %s is yet used. Not proceeding.", name) + continue + logger.debug( + '"%s" being pushed onto "%s"...', cmd["name"], lifopath + ) + PLifo.Push(lifopath, cmd["name"]) + logger.debug("Client %s has been newly created.", cmd["name"]) + chan = client_chans.get(cmd["name"], None) logger.debug("Starting new client in a thread...") client = threading.Thread( name=cmd["name"], target=functools.partial( - listen_client, name=cmd["name"], chan=q, cmd_map=cmd_map + listen_client, name=cmd["name"], chan=chan, cmd_map=cmd_map ), ) client.start() logger.debug('Client "%s" started.', cmd["name"]) continue - q.put(cmd) + + if cmd["cmd"] == "stop-daemon": + stop_all(client_chans) + logger.info("Daemon clients stopped. Exiting...") + raise StopDaemon + + logger.debug("Opening client: %s", cmd["name"]) + chan = client_chans.get(cmd["name"], None) + if chan is None: + logger.warning( + 'No client found with name "%s". Skipping.', cmd["name"] + ) + continue + + chan.put(cmd) + + +def stop_all(client_chans: T.Dict[str, queue.Queue]) -> None: + """ + Send the stop command to all the clients. + """ + for name, chan in client_chans.items(): + logger.debug('Stopping "%s"...', name) + chan.put({"name": name, "cmd": "stop"}) + logger.info("Stopped %s.", name)