#!/usr/bin/env python2.7
#
# Sample snippet showing threads and process control, purpose is to keep the
# process always alive.
#
# Rick van der Zwet <info@rickvanderzwet.nl>
#

import logging
import os
import subprocess
import threading
import time

# TODO(rvdz) Use package enum to generate this list:
# https://pypi.python.org/pypi/enum
class SM:
   ''' State definitions'''
   START, POLL, RESTART, SHUTDOWN = range(1,5)

class DEFAULT:
  cmd_logfile = 'file://stdout'

class RunCmd(threading.Thread):
  '''
  Create external command and verify process keeps running, the command log if
  writting to a file if needed
  '''
  def __init__(self, name, cmd, cmd_logfile=DEFAULT.cmd_logfile, restart_timeout=1):
    threading.Thread.__init__(self)
    self.name = name
    self.cmd = cmd
    self.cmd_logfile = cmd_logfile
    self.restart_timeout = restart_timeout


    # Debug logging and command logging
    self.logger = logging.getLogger('thread.' + name)
    if cmd_logfile == DEFAULT.cmd_logfile:
      self.cmd_log_fh = subprocess.PIPE
    else:
      self.cmd_log_fh = open(self.cmd_logfile, 'a')

    # Internal variables
    self.do_shutdown = False
    self.p = None
    self.start_time = 0
    self.state = SM.START

  def is_alive(self):
    ''' Shortcut function '''
    return (self.p.poll() == None)

  def safe_sleep(self, sleeptime):
    ''' Avoid sleeping for long periods and not beeing able to handle any signals'''
    while sleeptime > 0 and not self.do_shutdown:
      sleeptime -= 1
      time.sleep(1)

  def run(self):
    self.logger.info('Starting process "%s"', self.cmd)
    self.logger.info('logging command output to %s', self.cmd_logfile)
    # Run in infinite loop to allow restaring application, programmed using a
    # state-machine idea
    while True:
      if self.state == SM.START:
        self.start_time = int(time.time())
        self.p = subprocess.Popen(
          self.cmd,
          shell=True,
          stdin=None,
          stderr=subprocess.STDOUT,
          stdout=self.cmd_log_fh,
        )
        # Give process some time to start
        time.sleep(1)
        self.logger.info("Started process with pid: %i", self.p.pid)
        self.state = SM.POLL
      elif self.state == SM.POLL:
        if self.is_alive():
          # Sleep a little while to avoid polling non-stop.
          time.sleep(1)
        else:
	  runtime = int(time.time()) - self.start_time
          self.logger.warning("Process has stopped after %s seconds", runtime)
          self.state = SM.RESTART
        # Check whether a external shutdown is triggered
        if self.do_shutdown:
          self.state = SM.SHUTDOWN
      elif self.state == SM.RESTART:
        # Wait a little while to avoid re-spawning the process all the time.
        self.logger.info('re-starting process in %s second(s)', self.restart_timeout)
        self.safe_sleep(self.restart_timeout)
        self.state = SM.START
      elif self.state == SM.SHUTDOWN:
        if self.is_alive():
          self.logger.debug('Sending SIGTERM to pid %i', self.p.pid) 
          self.p.terminate()
          time.sleep(2)
        if self.is_alive():
          self.logger.debug('Not responding to SIGTERM sending SIGKILL to %i', self.p.pid)
          self.p.kill()
          self.p.wait()
        # Housekeeping on shutdown  
        if self.cmd_logfile != DEFAULT.cmd_logfile:
           self.cmd_log_fh.close()
        return
      else:
        self.logger.critical("Should never be able to reach state: %i", self.state)
        self.state = SM.SHUTDOWN

  def stop(self):
    self.logger.debug('Received shutdown trigger!')
    self.do_shutdown = True



if __name__ == '__main__':
  logging.basicConfig(level=logging.DEBUG)
  logger = logging.getLogger()
  
  # Configure some sample threads
  threads = {}
  threads['first'] = RunCmd('first', 'date && sleep 10')
  threads['second'] = RunCmd('second', 'date && sleep 5')
  
  # Starting all threads
  for _,thread in threads.items():
    thread.start()
  logger.info('The main program runs in the foreground')
  
  try:
    # Keep in foreground, mind not to use join and such as they do not work with
    # interrupt handling: http://bugs.python.org/issue1167930#msg56947
    while True:
      time.sleep(60)
  except (KeyboardInterrupt, SystemExit) as e:
    logger.info('Received signal to exit (%s), shutting down threads' % e)
    # Could combine this loops, but I rather have all my threads to close as
    # soon as possible.
    for _,thread in threads.items():
      logger.info("Signalling thread %s to shutdown", thread.name)
      thread.stop()
    for _,thread in threads.items():
      thread.join()
      logger.info("Thread %s has stopped", thread.name)
    logger.info('All done, goodbye!')