content: #  Copyright (c) ZenML GmbH 2021. All Rights Reserved.
  #
  #  Licensed under the Apache License, Version 2.0 (the "License");
  #  you may not use this file except in compliance with the License.
  #  You may obtain a copy of the License at:
  #
  #       https://www.apache.org/licenses/LICENSE-2.0
  #
  #  Unless required by applicable law or agreed to in writing, software
  #  distributed under the License is distributed on an "AS IS" BASIS,
  #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
  #  or implied. See the License for the specific language governing
  #  permissions and limitations under the License.
  """Utility functions to start/stop daemon processes.
  
  This is only implemented for UNIX systems and therefore doesn't work on
  Windows. Based on
  https://www.jejik.com/articles/2007/02/a_simple_unix_linux_daemon_in_python/
  """
  
  import atexit
  import io
  import os
  import signal
  import sys
  import types
  from typing import Any, Callable, Optional, TypeVar, cast
  
  import psutil
  
  from zenml.logger import get_logger
  
  logger = get_logger(__name__)
  
  # TODO [ENG-235]: Investigate supporting Windows if Windows can run Kubeflow.
  
  
  F = TypeVar("F", bound=Callable[..., Any])
  
  CHILD_PROCESS_WAIT_TIMEOUT = 5
  
  
  def daemonize(
      pid_file: str,
      log_file: Optional[str] = None,
      working_directory: str = "/",
  ) -> Callable[[F], F]:
      """Decorator that executes the decorated function as a daemon process.
  
      Use this decorator to easily transform any function into a daemon
      process.
  
      For example,
  
      ```python
      import time
      from zenml.utils.daemon import daemonize
  
  
      @daemonize(log_file='/tmp/daemon.log', pid_file='/tmp/daemon.pid')
      def sleeping_daemon(period: int) -> None:
          print(f"I'm a daemon! I will sleep for {period} seconds.")
          time.sleep(period)
          print("Done sleeping, flying away.")
  
      sleeping_daemon(period=30)
  
      print("I'm the daemon's parent!.")
      time.sleep(10) # just to prove that the daemon is running in parallel
      ```
  
      Args:
          pid_file: a file where the PID of the daemon process will
              be stored.
          log_file: file where stdout and stderr are redirected for the daemon
              process. If not supplied, the daemon will be silenced (i.e. have
              its stdout/stderr redirected to /dev/null).
          working_directory: working directory for the daemon process,
              defaults to the root directory.
  
      Returns:
          Decorated function that, when called, will detach from the current
          process and continue executing in the background, as a daemon
          process.
      """
  
      def inner_decorator(_func: F) -> F:
          def daemon(*args: Any, **kwargs: Any) -> None:
              """Standard daemonization of a process.
  
              Args:
                  *args: Arguments to be passed to the decorated function.
                  **kwargs: Keyword arguments to be passed to the decorated
                      function.
              """
              if sys.platform == "win32":
                  logger.error(
                      "Daemon functionality is currently not supported on Windows."
                  )
              else:
                  run_as_daemon(
                      _func,
                      log_file=log_file,
                      pid_file=pid_file,
                      working_directory=working_directory,
                      *args,
                      **kwargs,
                  )
  
          return cast(F, daemon)
  
      return inner_decorator
  
  
  def setup_daemon(
      pid_file: Optional[str] = None, log_file: Optional[str] = None
  ) -> None:
      """Sets up a daemon process.
  
      Args:
          pid_file: Path to file in which to store the PID of the daemon
              process.
          log_file: Optional file to which the daemons stdout/stderr will be
              redirected to.
      """
      # redirect standard file descriptors to devnull (or the given logfile)
      devnull = "/dev/null"
      if hasattr(os, "devnull"):
          devnull = os.devnull
  
      devnull_fd = os.open(devnull, os.O_RDWR)
      log_fd = (
          os.open(log_file, os.O_CREAT | os.O_RDWR | os.O_APPEND)
          if log_file
          else None
      )
      out_fd = log_fd or devnull_fd
  
      try:
          os.dup2(devnull_fd, sys.stdin.fileno())
      except io.UnsupportedOperation:
          # stdin is not a file descriptor
          pass
      try:
          os.dup2(out_fd, sys.stdout.fileno())
      except io.UnsupportedOperation:
          # stdout is not a file descriptor
          pass
      try:
          os.dup2(out_fd, sys.stderr.fileno())
      except io.UnsupportedOperation:
          # stderr is not a file descriptor
          pass
  
      # Close original FDs after duplication to avoid leaks.
      try:
          os.close(devnull_fd)
      except OSError:
          pass
  
      if log_fd is not None:
          try:
              os.close(log_fd)
          except OSError:
              pass
  
      if pid_file:
          # write the PID file
          with open(pid_file, "w+") as f:
              f.write(f"{os.getpid()}\n")
  
      # register actions in case this process exits/gets killed
      def cleanup() -> None:
          """Daemon cleanup."""
          sys.stderr.write("Cleanup: terminating children processes...\n")
          terminate_children()
          if pid_file and os.path.exists(pid_file):
              sys.stderr.write(f"Cleanup: removing PID file {pid_file}...\n")
              os.remove(pid_file)
          sys.stderr.flush()
  
      def sighndl(signum: int, frame: Optional[types.FrameType]) -> None:
          """Daemon signal handler.
  
          Args:
              signum: Signal number.
              frame: Frame object.
          """
          sys.stderr.write(f"Handling signal {signum}...\n")
          cleanup()
  
      signal.signal(signal.SIGTERM, sighndl)
      signal.signal(signal.SIGINT, sighndl)
      atexit.register(cleanup)
  
  
  def stop_process(pid: int) -> None:
      """Stops a process.
  
      Args:
          pid: PID of the process to stop.
      """
      if not psutil.pid_exists(pid):
          logger.debug("PID %s does not exist.", pid)
          return
  
      process = psutil.Process(pid)
      process.terminate()
  
      try:
          process.wait(CHILD_PROCESS_WAIT_TIMEOUT)
      except psutil.TimeoutExpired:
          logger.debug("Daemon PID %s did not terminate in time; killing.", pid)
          try:
              process.kill()
              try:
                  process.wait(CHILD_PROCESS_WAIT_TIMEOUT)
              except psutil.TimeoutExpired:
                  logger.error("Failed to kill daemon PID %s.", pid)
          except Exception as e:
              logger.error("Failed to kill daemon PID %s: %s", pid, e)
  
  
  def stop_daemon(pid_file: str) -> None:
      """Stops a daemon process.
  
      Args:
          pid_file: Path to file containing the PID of the daemon process to
              kill.
      """
      try:
          with open(pid_file, "r") as f:
              pid = int(f.read().strip())
      except (IOError, FileNotFoundError):
          logger.debug("Daemon PID file '%s' does not exist.", pid_file)
          return
  
      stop_process(pid)
  
  
  def get_daemon_pid_if_running(pid_file: str) -> Optional[int]:
      """Read and return the PID value from a PID file.
  
      It does this if the daemon process tracked by the PID file is running.
  
      Args:
          pid_file: Path to file containing the PID of the daemon
              process to check.
  
      Returns:
          The PID of the daemon process if it is running, otherwise None.
      """
      try:
          with open(pid_file, "r") as f:
              pid = int(f.read().strip())
      except (IOError, FileNotFoundError):
          logger.debug(
              f"Daemon PID file '{pid_file}' does not exist or cannot be read.",
          )
          return None
      except ValueError:
          logger.debug(f"Daemon PID file '{pid_file}' contains invalid data.")
          return None
  
      if not pid or not psutil.pid_exists(pid):
          logger.debug(f"Daemon with PID '{pid}' is no longer running.")
          return None
  
      logger.debug(f"Daemon with PID '{pid}' is running.")
      return pid
  
  
  def check_if_daemon_is_running(pid_file: str) -> bool:
      """Checks whether a daemon process indicated by the PID file is running.
  
      Args:
          pid_file: Path to file containing the PID of the daemon
              process to check.
  
      Returns:
          True if the daemon process is running, otherwise False.
      """
      return get_daemon_pid_if_running(pid_file) is not None
  
  
  def terminate_children() -> None:
      """Terminate all processes that are children of the currently running process."""
      pid = os.getpid()
      try:
          parent = psutil.Process(pid)
      except psutil.Error:
          # could not find parent process id
          return
      children = parent.children(recursive=False)
  
      for p in children:
          sys.stderr.write(f"Terminating child process with PID {p.pid}...\n")
          p.terminate()
      _, alive = psutil.wait_procs(children, timeout=CHILD_PROCESS_WAIT_TIMEOUT)
      for p in alive:
          sys.stderr.write(f"Killing child process with PID {p.pid}...\n")
          p.kill()
      _, alive = psutil.wait_procs(children, timeout=CHILD_PROCESS_WAIT_TIMEOUT)
  
  
  if sys.platform == "win32":
      logger.warning(
          "Daemon functionality is currently not supported on Windows."
      )
  else:
  
      def run_as_daemon(
          daemon_function: F,
          *args: Any,
          pid_file: str,
          log_file: Optional[str] = None,
          working_directory: str = "/",
          **kwargs: Any,
      ) -> None:
          """Runs a function as a daemon process.
  
          Args:
              daemon_function: The function to run as a daemon.
              pid_file: Path to file in which to store the PID of the daemon
                  process.
              log_file: Optional file to which the daemons stdout/stderr will be
                  redirected to.
              working_directory: Working directory for the daemon process,
                  defaults to the root directory.
              args: Positional arguments to pass to the daemon function.
              kwargs: Keyword arguments to pass to the daemon function.
  
          Raises:
              FileExistsError: If the PID file already exists.
          """
          # convert to absolute path as we will change working directory later
          if pid_file:
              pid_file = os.path.abspath(pid_file)
          if log_file:
              log_file = os.path.abspath(log_file)
  
          # create parent directory if necessary
          for f in (pid_file, log_file):
              if f:
                  os.makedirs(os.path.dirname(f), exist_ok=True)
  
          # check if PID file exists
          if pid_file and os.path.exists(pid_file):
              pid = get_daemon_pid_if_running(pid_file)
              if pid:
                  raise FileExistsError(
                      f"The PID file '{pid_file}' already exists and a daemon "
                      f"process with the same PID '{pid}' is already running."
                      f"Please remove the PID file or kill the daemon process "
                      f"before starting a new daemon."
                  )
              logger.warning(
                  f"Removing left over PID file '{pid_file}' from a previous "
                  f"daemon process that didn't shut down correctly."
              )
              os.remove(pid_file)
  
          # first fork
          try:
              pid = os.fork()
              if pid > 0:
                  # this is the process that called `run_as_daemon` so we
                  # wait for the child process to finish to avoid creating
                  # zombie processes. Then we simply return so the current process
                  # can continue what it was doing.
                  os.wait()
                  return
          except OSError as e:
              logger.error("Unable to fork (error code: %d)", e.errno)
              sys.exit(1)
  
          # decouple from parent environment
          os.chdir(working_directory)
          os.setsid()
          os.umask(0o22)
  
          # second fork
          try:
              pid = os.fork()
              if pid > 0:
                  # this is the parent of the future daemon process, kill it
                  # so the daemon gets adopted by the init process.
                  # we use os._exit here to prevent the inherited code from
                  # catching the SystemExit exception and doing something else.
                  os._exit(0)
          except OSError as e:
              sys.stderr.write(f"Unable to fork (error code: {e.errno})")
              # we use os._exit here to prevent the inherited code from
              # catching the SystemExit exception and doing something else.
              os._exit(1)
  
          # setup the daemon
          setup_daemon(pid_file, log_file)
  
          # finally run the actual daemon code
          daemon_function(*args, **kwargs)
          sys.exit(0)
  