recipe_modules/daemonizer/resources/daemonizer.py

#!/usr/bin/python
# Copyright 2017 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""
This script can daemonize another script. This is useful in running
background processes from recipes.
"""

import argparse
import logging
import os
import signal
import subprocess
import sys


def daemonize(cmd, pidfile):
    """
    This function is based on the Python recipe provided here:
    http://www.jejik.com/articles/2007/02/a_simple_unix_linux_daemon_in_python/
    """
    # flush the output
    sys.stdout.flush()
    sys.stderr.flush()

    # spawn a detached child process
    try:
        pid = os.fork()
        if pid > 0:
            # exit first parent
            sys.exit(0)
    except OSError, e:
        sys.stderr.write(
            "fork #1 failed, unable to daemonize: %d (%s)\n" % (e.errno, e.strerror)
        )
        sys.exit(1)

    # decouple from parent environment
    os.chdir("/")
    os.setsid()
    os.umask(0)

    # do second fork
    try:
        pid = os.fork()
        if pid > 0:
            # exit from second parent
            sys.exit(0)
    except OSError, e:
        sys.stderr.write(
            "fork #2 failed, unable to daemonize: %d (%s)\n" % (e.errno, e.strerror)
        )
        sys.exit(1)

    # redirect standard file descriptors
    sys.stdout.flush()
    sys.stderr.flush()
    si = file("/dev/null", "r")
    so = file("/dev/null", "a+")
    se = file("/dev/null", "a+", 0)
    os.dup2(si.fileno(), sys.stdin.fileno())
    os.dup2(so.fileno(), sys.stdout.fileno())
    os.dup2(se.fileno(), sys.stderr.fileno())

    proc = subprocess.Popen(cmd)

    # Write pid to file if applicable.
    if pidfile:
        try:
            with open(pidfile, "w") as pid_file:
                pid_file.write("%s" % str(proc.pid))
        except (IOError):
            logging.exception("Unable to write pid to file")

    proc.communicate()
    return proc.returncode


def stop(pidfile):
    if not pidfile:
        logging.error("pidfile arg must be specified when stopping a daemon")
        return 1
    try:
        with open(pidfile) as pid_file:
            pid = int(pid_file.readline())
        logging.info("Sending SIGTERM to %d", pid)
        os.kill(pid, signal.SIGTERM)
        os.remove(pidfile)
    except (IOError, OSError):
        logging.exception("Error terminating daemon process")


def restart(cmd, pidfile):
    # check for the pidfile to see if the daemon's already running
    if not pidfile:
        logging.error("pidfile arg must be specified when restarting a daemon")
        return 1
    try:
        with open(pidfile, "r") as pid_file:
            pid = int(pid_file.readline())
    except (IOError, ValueError):
        pid = None

    if pid:
        logging.info(
            "%s pid file already exists, attempting to kill process %d", pidfile, pid,
        )
        try:
            os.kill(pid, signal.SIGTERM)
        except OSError:
            logging.exception("Unable to kill old daemon process")

    return daemonize(cmd, pidfile)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--action",
        default="daemonize",
        choices=["daemonize", "stop", "restart"],
        help="What action to take",
    )
    parser.add_argument(
        "--pidfile", type=str, help="Path to store the daemon's pid",
    )
    parser.add_argument("cmd", help="Command to daemonize", nargs=argparse.REMAINDER)
    args = parser.parse_args()

    if args.action == "daemonize":
        return daemonize(args.cmd, args.pidfile)
    elif args.action == "stop":
        return stop(args.pidfile)
    elif args.action == "restart":
        return restart(args.cmd, args.pidfile)


if __name__ == "__main__":
    sys.exit(main())