from future.builtins import object
import psutil
import os
import time
import signal
import shlex
import gevent
try:
import subprocess32 as subprocess
except:
import subprocess
from .context import log
class Process(object):
""" The parent class of Worker, Agent and Supervisor """
exitcode = 0
def install_signal_handlers(self):
""" Handle events like Ctrl-C from the command line. """
self.graceful_stop = False
def request_shutdown_now():
self.shutdown_now()
def request_shutdown_graceful():
# Second time CTRL-C, shutdown now
if self.graceful_stop:
self.shutdown_now()
else:
self.graceful_stop = True
self.shutdown_graceful()
# First time CTRL-C, try to shutdown gracefully
gevent.signal_handler(signal.SIGINT, request_shutdown_graceful)
# User (or Heroku) requests a stop now, just mark tasks as interrupted.
gevent.signal_handler(signal.SIGTERM, request_shutdown_now)
class ProcessPool(object):
""" Manages a pool of processes """
def __init__(self, watch_interval=1, extra_env=None):
self.processes = []
self.desired_commands = []
self.greenlet_watch = None
self.watch_interval = watch_interval
self.stopping = False
self.extra_env = extra_env
def set_commands(self, commands, timeout=None):
""" Sets the processes' desired commands for this pool and manages diff to reach that state """
self.desired_commands = commands
target_commands = list(self.desired_commands)
for process in list(self.processes):
found = False
for i in range(len(target_commands)):
if process["command"] == target_commands[i]:
target_commands.pop(i)
found = True
break
if not found:
self.stop_process(process, timeout)
# What is left are the commands to add
# TODO: we should only do this once memory conditions allow
for command in target_commands:
self.spawn(command)
def spawn(self, command):
""" Spawns a new process and adds it to the pool """
# process_name
# output
# time before starting (wait for port?)
# start_new_session=True : avoid sending parent signals to child
env = dict(os.environ)
env["MRQ_IS_SUBPROCESS"] = "1"
env.update(self.extra_env or {})
# Extract env variables from shell commands.
parts = shlex.split(command)
for p in list(parts):
if "=" in p:
env[p.split("=")[0]] = p[len(p.split("=")[0]) + 1:]
parts.pop(0)
else:
break
p = subprocess.Popen(parts, shell=False, close_fds=True, env=env, cwd=os.getcwd())
self.processes.append({
"subprocess": p,
"pid": p.pid,
"command": command,
"psutil": psutil.Process(pid=p.pid)
})
def start(self):
self.greenlet_watch = gevent.spawn(self.watch)
self.greenlet_watch.start()
def wait(self):
""" Waits for the pool to be fully stopped """
while True:
if not self.greenlet_watch:
break
if self.stopping:
gevent.sleep(0.1)
else:
gevent.sleep(1)
def watch(self):
while True:
self.watch_processes()
gevent.sleep(self.watch_interval)
def watch_processes(self):
""" Manages the status of all the known processes """
for process in list(self.processes):
self.watch_process(process)
# Cleanup processes
self.processes = [p for p in self.processes if not p.get("dead")]
if self.stopping and len(self.processes) == 0:
self.stop_watch()
def watch_process(self, process):
""" Manages the status of a single process """
status = process["psutil"].status()
# TODO: how to avoid zombies?
# print process["pid"], status
if process.get("terminate"):
if status in ("zombie", "dead"):
process["dead"] = True
elif process.get("terminate_at"):
if time.time() > (process["terminate_at"] + 5):
log.warning("Process %s had to be sent SIGKILL" % (process["pid"], ))
process["subprocess"].send_signal(signal.SIGKILL)
elif time.time() > process["terminate_at"]:
log.warning("Process %s had to be sent SIGTERM" % (process["pid"], ))
process["subprocess"].send_signal(signal.SIGTERM)
else:
if status in ("zombie", "dead"):
# Restart a new process right away (TODO: sleep a bit? max retries?)
process["dead"] = True
self.spawn(process["command"])
elif status not in ("running", "sleeping"):
log.warning("Process %s was in status %s" % (process["pid"], status))
# process["subprocess"].returncode in (0, 2, 3)
def stop(self, timeout=None):
""" Initiates a graceful stop of the processes """
self.stopping = True
for process in list(self.processes):
self.stop_process(process, timeout=timeout)
def stop_process(self, process, timeout=None):
""" Initiates a graceful stop of one process """
process["terminate"] = True
if timeout is not None:
process["terminate_at"] = time.time() + timeout
process["subprocess"].send_signal(signal.SIGINT)
def terminate(self):
""" Terminates the processes right now with a SIGTERM """
for process in list(self.processes):
process["subprocess"].send_signal(signal.SIGTERM)
self.stop_watch()
def kill(self):
""" Kills the processes right now with a SIGKILL """
for process in list(self.processes):
process["subprocess"].send_signal(signal.SIGKILL)
self.stop_watch()
def stop_watch(self):
""" Stops the periodic watch greenlet, thus the pool itself """
if self.greenlet_watch:
self.greenlet_watch.kill(block=False)
self.greenlet_watch = None
