xref: /openbmc/openbmc-test-automation/lib/func_timer.py (revision 20f38712b324e61a94e174017c487a0af4b373e1)

#!/usr/bin/env python3

r"""
Define the func_timer class.
"""

import os
import signal
import sys
import time

import gen_misc as gm
import gen_print as gp
import gen_valid as gv


class func_timer_class:
    r"""
    Define the func timer class.

    A func timer object can be used to run any function/arguments but with an additional benefit of being
    able to specify a time_out value.  If the function fails to complete before the timer expires, a
    ValueError exception will be raised along with a detailed error message.

    Example code:

    func_timer = func_timer_class()
    func_timer.run(run_key, "sleep 2", time_out=1)

    In this example, the run_key function is being run by the func_timer object with a time_out value of 1
    second.  "sleep 2" is a positional parm for the run_key function.
    """

    def __init__(self, obj_name="func_timer_class"):
        # Initialize object variables.
        self.__obj_name = obj_name
        self.__func = None
        self.__time_out = None
        self.__child_pid = 0
        # Save the original SIGUSR1 handler for later restoration by this class' methods.
        self.__original_SIGUSR1_handler = signal.getsignal(signal.SIGUSR1)

    def __del__(self):
        self.cleanup()

    def sprint_obj(self):
        r"""
        sprint the fields of this object.  This would normally be for debug purposes.
        """

        buffer = ""
        buffer += self.__class__.__name__ + ":\n"
        indent = 2
        try:
            func_name = self.__func.__name__
        except AttributeError:
            func_name = ""
        buffer += gp.sprint_var(func_name, indent=indent)
        buffer += gp.sprint_varx("time_out", self.__time_out, indent=indent)
        buffer += gp.sprint_varx("child_pid", self.__child_pid, indent=indent)
        buffer += gp.sprint_varx(
            "original_SIGUSR1_handler",
            self.__original_SIGUSR1_handler,
            indent=indent,
        )
        return buffer

    def print_obj(self):
        r"""
        print the fields of this object to stdout.  This would normally be for debug purposes.
        """

        sys.stdout.write(self.sprint_obj())

    def cleanup(self):
        r"""
        Cleanup after the run method.
        """

        try:
            gp.lprint_executing()
            gp.lprint_var(self.__child_pid)
        except (AttributeError, KeyError, TypeError):
            # NOTE: In python 3, this code fails with "KeyError: ('__main__',)" when calling functions like
            # lprint_executing that use inspect.stack() during object destruction.  No fixes found so
            # tolerating the error.  In python 2.x, it may fail with TypeError.  This seems to happen when
            # cleaning up after an exception was raised.
            pass

        # If self.__child_pid is 0, then we are either running as the child or we've already cleaned up.
        # If self.__time_out is None, then no child process would have been spawned.
        if self.__child_pid == 0 or self.__time_out is None:
            return

        # Restore the original SIGUSR1 handler.
        if self.__original_SIGUSR1_handler != 0:
            signal.signal(signal.SIGUSR1, self.__original_SIGUSR1_handler)
        try:
            gp.lprint_timen("Killing child pid " + str(self.__child_pid) + ".")
            os.kill(self.__child_pid, signal.SIGKILL)
        except OSError:
            gp.lprint_timen("Tolerated kill failure.")
        try:
            gp.lprint_timen("os.waitpid(" + str(self.__child_pid) + ")")
            os.waitpid(self.__child_pid, 0)
        except OSError:
            gp.lprint_timen("Tolerated waitpid failure.")
        self.__child_pid = 0
        # For debug purposes, prove that the child pid was killed.
        children = gm.get_child_pids()
        gp.lprint_var(children)

    def timed_out(self, signal_number, frame):
        r"""
        Handle a SIGUSR1 generated by the child process after the time_out has expired.

        signal_number               The signal_number of the signal causing this method to get invoked.  This
                                    should always be 10 (SIGUSR1).
        frame                       The stack frame associated with the function that times out.
        """

        gp.lprint_executing()

        self.cleanup()

        # Compose an error message.
        err_msg = "The " + self.__func.__name__
        err_msg += " function timed out after " + str(self.__time_out)
        err_msg += " seconds.\n"
        if not gp.robot_env:
            err_msg += gp.sprint_call_stack()

        raise ValueError(err_msg)

    def run(self, func, *args, **kwargs):
        r"""
        Run the indicated function with the given args and kwargs and return the value that the function
        returns.  If the time_out value expires, raise a ValueError exception with a detailed error message.

        This method passes all of the args and kwargs directly to the child function with the following
        important exception: If kwargs contains a 'time_out' value, it will be used to set the func timer
        object's time_out value and then the kwargs['time_out'] entry will be removed.  If the time-out
        expires before the function finishes running, this method will raise a ValueError.

        Example:
        func_timer = func_timer_class()
        func_timer.run(run_key, "sleep 3", time_out=2)

        Example:
        try:
            result = func_timer.run(func1, "parm1", time_out=2)
            print_var(result)
        except ValueError:
            print("The func timed out but we're handling it.")

        Description of argument(s):
        func                        The function object which is to be called.
        args                        The arguments which are to be passed to the function object.
        kwargs                      The keyword arguments which are to be passed to the function object.  As
                                    noted above, kwargs['time_out'] will get special treatment.
        """

        gp.lprint_executing()

        # Store method parms as object parms.
        self.__func = func

        # Get self.__time_out value from kwargs.  If kwargs['time_out'] is not present, self.__time_out will
        # default to None.
        self.__time_out = None
        if "time_out" in kwargs:
            self.__time_out = kwargs["time_out"]
            del kwargs["time_out"]
            # Convert "none" string to None.
            try:
                if self.__time_out.lower() == "none":
                    self.__time_out = None
            except AttributeError:
                pass
            if self.__time_out is not None:
                self.__time_out = int(self.__time_out)
                # Ensure that time_out is non-negative.
                message = gv.valid_range(
                    self.__time_out, 0, var_name="time_out"
                )
                if message != "":
                    raise ValueError(
                        "\n" + gp.sprint_error_report(message, format="long")
                    )

        gp.lprint_varx("time_out", self.__time_out)
        self.__child_pid = 0
        if self.__time_out is not None:
            # Save the original SIGUSR1 handler for later restoration by this class' methods.
            self.__original_SIGUSR1_handler = signal.getsignal(signal.SIGUSR1)
            # Designate a SIGUSR1 handling function.
            signal.signal(signal.SIGUSR1, self.timed_out)
            parent_pid = os.getpid()
            self.__child_pid = os.fork()
            if self.__child_pid == 0:
                gp.dprint_timen(
                    "Child timer pid "
                    + str(os.getpid())
                    + ": Sleeping for "
                    + str(self.__time_out)
                    + " seconds."
                )
                time.sleep(self.__time_out)
                gp.dprint_timen(
                    "Child timer pid "
                    + str(os.getpid())
                    + ": Sending SIGUSR1 to parent pid "
                    + str(parent_pid)
                    + "."
                )
                os.kill(parent_pid, signal.SIGUSR1)
                os._exit(0)

        # Call the user's function with the user's arguments.
        children = gm.get_child_pids()
        gp.lprint_var(children)
        gp.lprint_timen("Calling the user's function.")
        gp.lprint_varx("func_name", func.__name__)
        gp.lprint_vars(args, kwargs)
        try:
            result = func(*args, **kwargs)
        except Exception as func_exception:
            # We must handle all exceptions so that we have the chance to cleanup before re-raising the
            # exception.
            gp.lprint_timen("Encountered exception in user's function.")
            self.cleanup()
            raise (func_exception)
        gp.lprint_timen("Returned from the user's function.")

        self.cleanup()

        return result