Revised documentation and coding style in daemonizer.py library.

pyzenkit/daemonizer.py

-#!/usr/bin/python
+#!/usr/bin/env python3
 # -*- coding: utf-8 -*-
 #-------------------------------------------------------------------------------
 # Copyright (C) since 2016 Jan Mach <honza.mach.ml@gmail.com>
@@ -6,16 +6,39 @@
 # Use of this source is governed by the MIT license, see LICENSE file.
 #-------------------------------------------------------------------------------
 
+
 """
-Daemonization library.
+This module contains simple daemonization library, that takes care of all tasks
+necessary to correctly daemonize a process. Correct daemonization consists of
+following steps:
+
+* Setup directories and limits
+* Setup user and group permissions
+* Double fork and split session
+* Setup signal handlers
+* Close all open file descriptors (except for possible log files)
+* Redirect ``stdin``, ``stdout``, ``stderr`` to ``/dev/null``
+* Detect current PID and store it to appropriate PID file
+* At exit remove PID file
+
+These steps are performed during *full* daemonization. For many purposes it is however
+usefull to be capable of some kind of *lite* daemonization, in which case almost
+every step in previous list is done except for double forking, closing all files
+and redirecting ``std*`` to ``/dev/null``. This can be usefull during testing or debugging,
+or even during production deployments for some kind of *stay in foreground* feature, which
+still enables the users to control the application from outside with signals.
 """
 
+
+__author__  = "Jan Mach <honza.mach.ml@gmail.com>"
+__credits__ = "Pavel Kácha <ph@rook.cz>"
+
+
 import os
-import re
 import signal
-import resource
 import atexit
 
+
 def get_logger_files(logger):
     """
     Return file handlers of all currently active loggers.
@@ -25,8 +48,12 @@ def get_logger_files(logger):
 
     .. warning::
 
-        This method is hacking internal structure of different module and might
+        This method is hacking internal structure of external module and might
         stop working, although the related interface has been stable for a long time.
+
+    :param logging.Logger logger: Logger to be analyzed for open file descriptors.
+    :return: List of open file descriptors used by logging service.
+    :rtype: list
     """
     files = []
     for handler in logger.handlers:
@@ -36,75 +63,142 @@ def get_logger_files(logger):
             files.append(handler.socket)
     return files
 
-def write_pid(pidfile, pid):
+def write_pid(pid_file, pid):
     """
     Write given PID into given PID file.
+
+    :param str pidfile: Name of the PID file to write to.
+    :param int pid: PID to write.
     """
     if not isinstance(pid, int):
         raise Exception("Process PID must be integer")
-    pidfd = os.open(pidfile, os.O_RDWR|os.O_CREAT|os.O_EXCL|os.O_TRUNC)
+    pidfd = os.open(pid_file, os.O_RDWR|os.O_CREAT|os.O_EXCL|os.O_TRUNC)
     os.write(pidfd, bytes(str(pid)+"\n", 'UTF-8'))
     os.close(pidfd)
 
-def read_pid(pidfile):
+def read_pid(pid_file):
     """
     Read PID from given PID file.
+
+    :param str pidfile: Name of the PID file to read from.
+    :return: PID from given PID file.
+    :rtype: int
     """
-    with open(pidfile, 'r') as pidfd:
+    with open(pid_file, 'r') as pidfd:
         return int(pidfd.readline().strip())
 
-def daemonize_lite(
-        chroot_dir = None, work_dir = None, umask = None, uid = None, gid = None,
-        pidfile = None, signals = {}):
+
+#-------------------------------------------------------------------------------
+
+def _setup_fs(chroot_dir, work_dir, umask):
     """
-    Perform lite daemonization of currently running process.
+    Internal helper method, setup filesystem related features.
 
-    The lite daemonization does everything full daemonization does but detaching
-    from current session. This can be usefull when debugging daemons, because they
-    can be tested, benchmarked and profiled more easily.
+    :param str chroot_dir: Name of the chroot directory (may be ``None``).
+    :param str work_dir: Name of the work directory (may be ``None``).
+    :param int umask: Umask as octal number (eg. ``0o002`` or ``0o022``, may be ``None``).
     """
-    # Setup directories, limits, users, etc.
     if chroot_dir is not None:
         os.chdir(chroot_dir)
         os.chroot(chroot_dir)
-    if umask is not None:
-        os.umask(umask)
     if work_dir is not None:
         os.chdir(work_dir)
+    if umask is not None:
+        os.umask(umask)
+
+def _setup_perms(uid, gid):
+    """
+    Internal helper method, setup user and group permissions.
+
+    :param int uid: User ID to which to drop the permissions (may be ``None``).
+    :param int gid: Group ID to which to drop the permissions (may be ``None``).
+    """
     if gid is not None:
         os.setgid(gid)
     if uid is not None:
         os.setuid(uid)
 
-    # Setup signal handlers
-    for (signum, handler) in signals.items():
-        signal.signal(signum, handler)
+def _setup_sh(signals):
+    """
+    Internal helper method, setup desired signal handlers.
+
+    :param dict signals: Desired signal to be handled as keys and appropriate handlers as values (may be ``None``).
+    """
+    if signals is not None:
+        for (signum, handler) in signals.items():
+            signal.signal(signum, handler)
 
-    # Detect current process PID.
+def _setup_pf(pid_file):
+    """
+    Internal helper method, setup PID file and atexit cleanup callback.
+
+    :param str pid_file: Full path to the PID file (may be ``None``).
+    """
     pid = os.getpid()
 
-    # Create PID file and ensure its removal after current process is done.
-    if pidfile is not None:
-        if not pidfile.endswith('.pid'):
-            raise Exception("Invalid PID file name, it must end with '.pid' extension")
-        write_pid(pidfile, pid)
+    if pid_file is not None:
+        if not pid_file.endswith('.pid'):
+            raise ValueError("Invalid PID file name '{}', it must end with '.pid' extension".format(pid_file))
+
+        write_pid(pid_file, pid)
 
         # Define and setup 'atexit' closure, that will take care of removing pid file
         @atexit.register
         def unlink_pidfile():
+            """
+            Callback for removing PID file at application exit.
+            """
             try:
-                os.unlink(pidfile)
+                os.unlink(pid_file)
             except Exception:
                 pass
-        return (pid, pidfile)
+        return (pid, pid_file)
     else:
         return (pid, None)
 
+#-------------------------------------------------------------------------------
+
+
+def daemonize_lite(
+        chroot_dir = None, work_dir = None, umask = None, uid = None, gid = None,
+        pid_file = None, signals = None):
+    """
+    Perform lite daemonization of currently running process. All of the function
+    arguments are optinal, so that it is possible to easily turn on/off almost
+    any part of daemonization process. For example omitting the ``uid`` and ``gid``
+    arguments will result in process permissions not to be changed.
+
+    The lite daemonization does everything full daemonization does but detaching
+    from current session. This can be usefull when debugging daemons, because they
+    can be tested, benchmarked and profiled more easily.
+
+    :param str chroot_dir: Name of the chroot directory (may be ``None``).
+    :param str work_dir: Name of the work directory (may be ``None``).
+    :param int umask: Umask as octal number (eg. ``0o002`` or ``0o022``, may be ``None``).
+    :param int uid: User ID to which to drop the permissions (may be ``None``).
+    :param int gid: Group ID to which to drop the permissions (may be ``None``).
+    :param str pid_file: Full path to the PID file (may be ``None``).
+    :param dict signals: Desired signal to be handled as keys and appropriate handlers as values (may be ``None``).
+    """
+
+    # Setup directories, limits, users, etc.
+    _setup_fs(chroot_dir, work_dir, umask)
+    _setup_perms(uid, gid)
+
+    # Setup signal handlers.
+    _setup_sh(signals)
+
+    # Write PID into PID file.
+    return _setup_pf(pid_file)
+
 def daemonize(
         chroot_dir = None, work_dir = None, umask = None, uid = None, gid = None,
-        pidfile = None, files_preserve = [], signals = {}):
+        pid_file = None, files_preserve = None, signals = None):
     """
-    Perform full daemonization of currently running process.
+    Perform full daemonization of currently running process. All of the function
+    arguments are optinal, so that it is possible to easily turn on/off almost
+    any part of daemonization process. For example omitting the ``uid`` and ``gid``
+    arguments will result in process permissions not to be changed.
 
     NOTE: It would be possible to call daemonize_lite() method from within this
     method, howewer for readability purposes and to maintain correct ordering
@@ -112,32 +206,34 @@ def daemonize(
     two separate methods with similar contents. It will be necessary to update
     both when making any improvements, however I do not expect them to change
     much and often, if ever.
+
+    :param str chroot_dir: Name of the chroot directory (may be ``None``).
+    :param str work_dir: Name of the work directory (may be ``None``).
+    :param int umask: Umask as octal number (eg. ``0o002`` or ``0o022``, may be ``None``).
+    :param int uid: User ID to which to drop the permissions (may be ``None``).
+    :param int gid: Group ID to which to drop the permissions (may be ``None``).
+    :param str pid_file: Full path to the PID file (may be ``None``).
+    :param list files_preserve: List of file handles to preserve from closing (may be ``None``).
+    :param dict signals: Desired signal to be handled as keys and appropriate handlers as values (may be ``None``).
     """
+
     # Setup directories, limits, users, etc.
-    if chroot_dir is not None:
-        os.chdir(chroot_dir)
-        os.chroot(chroot_dir)
-    if umask is not None:
-        os.umask(umask)
-    if work_dir is not None:
-        os.chdir(work_dir)
-    if gid is not None:
-        os.setgid(gid)
-    if uid is not None:
-        os.setuid(uid)
+    _setup_fs(chroot_dir, work_dir, umask)
+    _setup_perms(uid, gid)
 
-    # Doublefork and split session.
+    # Doublefork and split session to fully detach from current terminal.
     if os.fork()>0:
         os._exit(0)
     os.setsid()
     if os.fork()>0:
         os._exit(0)
 
-    # Setup signal handlers
-    for (signum, handler) in signals.items():
-        signal.signal(signum, handler)
+    # Setup signal handlers.
+    _setup_sh(signals)
 
-    # Close all open file descriptors.
+    # Close all open file descriptors, except excluded files.
+    #if files_preserve is None:
+    #    files_preserve = []
     #descr_preserve = set(f.fileno() for f in files_preserve)
     #maxfd = resource.getrlimit(resource.RLIMIT_NOFILE)[1]
     #if maxfd==resource.RLIM_INFINITY:
@@ -151,61 +247,54 @@ def daemonize(
 
     # Redirect stdin, stdout, stderr to /dev/null.
     devnull = os.open(os.devnull, os.O_RDWR)
-    for fd in range(3):
-        os.dup2(devnull, fd)
+    for fdn in range(3):
+        os.dup2(devnull, fdn)
 
-    # Detect current process PID.
-    pid = os.getpid()
+    # Write PID into PID file.
+    return _setup_pf(pid_file)
 
-    # Create PID file and ensure its removal after current process is done.
-    if pidfile is not None:
-        if not pidfile.endswith('.pid'):
-            raise Exception("Invalid PID file name, it must end with '.pid' extension")
-        write_pid(pidfile, pid)
 
-        # Define and setup atexit closure
-        @atexit.register
-        def unlink_pidfile():
-            try:
-                os.unlink(pidfile)
-            except Exception:
-                pass
-        return (pid, pidfile)
-    else:
-        return (pid, None)
+#-------------------------------------------------------------------------------
 
+#
+# Perform the demonstration.
+#
 if __name__ == "__main__":
-    """
-    Perform the demonstration.
-    """
 
     def hnd_sig_hup(signum, frame):
-        print("Received signal HUP")
+        """Bogus handler for signal HUP for demonstration purposes."""
+        print("HANDLER CALLBACK: Received signal HUP ({})".format(signum))
+
     def hnd_sig_usr1(signum, frame):
-        print("Received signal USR1")
+        """Bogus handler for signal USR1 for demonstration purposes."""
+        print("HANDLER CALLBACK: Received signal USR1 ({})".format(signum))
+
     def hnd_sig_usr2(signum, frame):
-        print("Received signal USR2")
-
-    (pid, pidfile) = daemonize_lite(
-            work_dir = "/tmp",
-            pidfile = "/tmp/demo.pyzenkit.daemonizer.pid",
-            signals = {
-                signal.SIGHUP:  hnd_sig_hup,
-                signal.SIGUSR1: hnd_sig_usr1,
-                signal.SIGUSR2: hnd_sig_usr2,
-            }
-        )
+        """Bogus handler for signal USR2 for demonstration purposes."""
+        print("HANDLER CALLBACK: Received signal USR2 ({})".format(signum))
+
+    (PIDV, PIDF) = daemonize_lite(
+        work_dir = "/tmp",
+        pid_file = "/tmp/demo.pyzenkit.daemonizer.pid",
+        umask    = 0o022,
+        signals  = {
+            signal.SIGHUP:  hnd_sig_hup,
+            signal.SIGUSR1: hnd_sig_usr1,
+            signal.SIGUSR2: hnd_sig_usr2,
+        }
+    )
 
     print("Lite daemonization complete:")
-    print("\tPID: '{}'".format(pid))
-    print("\tPID file: '{}'".format(pidfile))
-    print("\tCWD: '{}'".format(os.getcwd()))
-    print("\tPID in PID file: '{}'".format(read_pid(pidfile)))
+    print("* PID:             '{}'".format(PIDV))
+    print("* PID file:        '{}'".format(PIDF))
+    print("* CWD:             '{}'".format(os.getcwd()))
+    print("* PID in PID file: '{}'".format(read_pid(PIDF)))
 
     print("Checking signal handling:")
-    os.kill(pid, signal.SIGHUP)
-    os.kill(pid, signal.SIGUSR1)
-    os.kill(pid, signal.SIGUSR2)
-    os.kill(read_pid(pidfile), signal.SIGHUP)
-    os.kill(read_pid(pidfile), signal.SIGUSR1)
-    os.kill(read_pid(pidfile), signal.SIGUSR2)
+    os.kill(PIDV, signal.SIGHUP)
+    os.kill(PIDV, signal.SIGUSR1)
+    os.kill(PIDV, signal.SIGUSR2)
+    print("Checking signal handling, read PID from PID file:")
+    os.kill(read_pid(PIDF), signal.SIGHUP)
+    os.kill(read_pid(PIDF), signal.SIGUSR1)
+    os.kill(read_pid(PIDF), signal.SIGUSR2)

pyzenkit/tests/test_daemonizer.py

@@ -6,6 +6,7 @@
 # Use of this source is governed by the MIT license, see LICENSE file.
 #-------------------------------------------------------------------------------
 
+
 import unittest
 from unittest.mock import Mock, MagicMock, call
 from pprint import pformat, pprint
@@ -15,20 +16,18 @@ import sys
 import shutil
 import signal
 
+
 # Generate the path to custom 'lib' directory
 lib = os.path.abspath(os.path.join(os.path.dirname(__file__), '../../'))
 sys.path.insert(0, lib)
 
 import pyzenkit.daemonizer
 
+
 PID_FILE = '/tmp/test.pyzenkit.daemonizer.pid'
 
-class TestPyzenkitDaemonizer(unittest.TestCase):
 
-    def setUp(self):
-        pass
-    def tearDown(self):
-        pass
+class TestPyzenkitDaemonizer(unittest.TestCase):
 
     def test_01_basic(self):
         """
@@ -49,26 +48,29 @@ class TestPyzenkitDaemonizer(unittest.TestCase):
         Perform lite daemonization tests.
         """
         def hnd_sig_hup(signum, frame):
-            print("Received signal HUP")
+            print("HANDLER CALLBACK: Received signal HUP ({})".format(signum))
+
         def hnd_sig_usr1(signum, frame):
-            print("Received signal USR1")
+            print("HANDLER CALLBACK: Received signal USR1 ({})".format(signum))
+
         def hnd_sig_usr2(signum, frame):
-            print("Received signal USR2")
+            print("HANDLER CALLBACK: Received signal USR2 ({})".format(signum))
 
         self.assertFalse(os.path.isfile(PID_FILE))
-        (pid, pidfile) = pyzenkit.daemonizer.daemonize_lite(
+        (pid, pid_file) = pyzenkit.daemonizer.daemonize_lite(
                 work_dir = '/tmp',
-                pidfile = PID_FILE,
-                signals = {
+                pid_file = PID_FILE,
+                umask    = 0o022,
+                signals  = {
                     signal.SIGHUP:  hnd_sig_hup,
                     signal.SIGUSR1: hnd_sig_usr1,
                     signal.SIGUSR2: hnd_sig_usr2,
                 },
             )
         self.assertTrue(os.path.isfile(PID_FILE))
-        self.assertTrue(os.path.isfile(pidfile))
+        self.assertTrue(os.path.isfile(pid_file))
         self.assertEqual(pyzenkit.daemonizer.read_pid(PID_FILE), pid)
-        self.assertEqual(pyzenkit.daemonizer.read_pid(pidfile), pid)
+        self.assertEqual(pyzenkit.daemonizer.read_pid(pid_file), pid)
         self.assertEqual(os.getcwd(), '/tmp')