""" Module for simple execution of external programs with keyword arguments. Also supports piping output into standard logging utilities. """ import logging import os import threading import sys import subprocess KWARG_TYPE_IGNORE = -1 KWARG_TYPE_NORMAL = 0 KWARG_TYPE_EQUAL = 1 KWARG_TYPE_MULTIPLE = 2 KWARG_TYPE_CALLBACK = 3 def apply_json_args(process, json_doc, custom_kwargs={}): """Translate keyword arguments (JSON) into an argument list for an external process. CALLBACK-type args can do arbitrary things to the process being started (set env vars, change the process name, etc.). """ for field in json_doc: kwarg, kwarg_type = ("--" + field, KWARG_TYPE_NORMAL) if field not in custom_kwargs \ else custom_kwargs[field][0:2] value = json_doc[field] if kwarg_type == KWARG_TYPE_NORMAL: if value is not None: process.arguments.append(kwarg) if str(value): process.arguments.append(str(value)) elif kwarg_type == KWARG_TYPE_EQUAL: process.arguments.append(kwarg + "=" + str(value)) elif kwarg_type == KWARG_TYPE_MULTIPLE: for ind_value in value: process.arguments.append(kwarg) process.arguments.append(str(ind_value)) elif kwarg_type == KWARG_TYPE_CALLBACK: cl_arg_callback = custom_kwargs[field][2] cl_arg_callback(process, field, value) class LoggerPipe(threading.Thread): """Monitors an external program's output and sends it to a logger.""" def __init__(self, logger, level, pipe_out): threading.Thread.__init__(self) self.logger = logger self.level = level self.pipe_out = pipe_out self.lock = threading.Lock() self.condition = threading.Condition(self.lock) self.started = False self.finished = False self.start() def run(self): with self.lock: self.started = True self.condition.notify_all() for line in self.pipe_out: self.logger.log(self.level, line.strip()) with self.lock: self.finished = True self.condition.notify_all() def wait_until_started(self): with self.lock: while not self.started: self.condition.wait() def wait_until_finished(self): with self.lock: while not self.finished: self.condition.wait() def flush(self): for handler in self.logger.handlers: handler.flush() class ExternalContext(object): def __init__(self, env=None, env_vars={}, logger=None, **kwargs): self.env = env self.env_vars = env_vars self.logger = logger if not logger: return logging.getLogger("") self.kwargs = dict(kwargs.items()) def clone(self): return ExternalContext(self.env, self.env_vars, self.logger, **self.kwargs) class ExternalProgram(object): """Encapsulates an execution of an external program. Unlike subprocess, does not immediately execute the program but allows for further configuration and setup. Converts keyword arguments in JSON into an argument list and allows for easy execution with custom environment variables. """ def __init__(self, executable, context=None, env=None, env_vars=None, custom_kwargs={}, **kwargs): self.executable = executable self.context = context if not self.context: self.context = ExternalContext(env, env_vars, **kwargs) else: self.context.kwargs.update(kwargs) self.custom_kwargs = custom_kwargs self.process = None def build_process(self, context=None): if not context: context = self.context process_kwargs = {} process_kwargs.update(context.kwargs) process = _Process(self.executable, env_vars=context.env_vars, logger=context.logger) apply_json_args(process, process_kwargs, self.custom_kwargs) return process def logger(self): return self.context.logger def start(self): self.process = self.build_process() self.process.start() def pid(self): return self.process.subprocess.pid def poll(self): return self.process.poll() def wait(self): return_code = self.process.wait() self.process = None return return_code def stop(self): return_code = self.process.stop() self.process = None return return_code def flush(self): self.process.flush() def __str__(self): return (self.process if self.process else self.build_process()).__str__() def __repr__(self): return self.__str__() class _Process(object): """The system-independent execution of an external program. Handles finicky stuff once we have our environment, arguments, and logger sorted out. """ def __init__(self, executable, arguments=[], env=None, env_vars=None, logger=None): self.executable = executable self.arguments = [] + arguments self.env = env self.env_vars = env_vars self.logger = logger self.subprocess = None self.stdout_logger = None self.stderr_logger = None # Windows only self.subprocess_job_object = None def start(self): argv, env = [self.executable] + self.arguments, self.env if self.env_vars: if not env: env = os.environ.copy() env.update(self.env_vars) creation_flags = 0 if os.sys.platform == "win32": # Magic number needed to allow job reassignment in Windows 7 # see: MSDN - Process Creation Flags - ms684863 CREATE_BREAKAWAY_FROM_JOB = 0x01000000 creation_flags = CREATE_BREAKAWAY_FROM_JOB stdout = sys.stdout if not self.logger else subprocess.PIPE stderr = sys.stderr if not self.logger else subprocess.PIPE self.subprocess = subprocess.Popen(argv, env=env, creationflags=creation_flags, stdout=stdout, stderr=stderr) if stdout == subprocess.PIPE: self.stdout_logger = LoggerPipe(self.logger, logging.INFO, self.subprocess.stdout) self.stdout_logger.wait_until_started() if stderr == subprocess.PIPE: self.stderr_logger = LoggerPipe(self.logger, logging.ERROR, self.subprocess.stderr) self.stderr_logger.wait_until_started() if os.sys.platform == "win32": # Create a job object with the "kill on job close" flag # This is inherited by child processes (i.e. the mongod started on our behalf by # buildlogger) and lets us terminate the whole tree of processes rather than # orphaning the mongod. import win32job job_object = win32job.CreateJobObject(None, '') job_info = win32job.QueryInformationJobObject( job_object, win32job.JobObjectExtendedLimitInformation) job_info['BasicLimitInformation']['LimitFlags'] |= \ win32job.JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE win32job.SetInformationJobObject(job_object, win32job.JobObjectExtendedLimitInformation, job_info) win32job.AssignProcessToJobObject(job_object, proc._handle) self.subprocess_job_object = job_object def poll(self): return self.subprocess.poll() def wait(self): return_code = self.subprocess.wait() self.flush() if self.stdout_logger: self.stdout_logger.wait_until_finished() self.stdout_logger = None if self.stderr_logger: self.stderr_logger.wait_until_finished() self.stderr_logger = None return return_code def stop(self): try: if os.sys.platform == "win32": import win32job win32job.TerminateJobObject(self.subprocess_job_object, -1) # Windows doesn't seem to kill the process immediately, so give # it some time to die time.sleep(5) elif hasattr(self.subprocess, "terminate"): # This method added in Python 2.6 self.subprocess.terminate() else: os.kill(self.subprocess.pid, 15) except Exception as e: print >> self.subprocess_outputs.stderr, "error shutting down process" print >> self.subprocess_outputs.stderr, e return self.wait() def flush(self): if self.subprocess: if not self.stderr_logger: # Going to the console sys.stderr.flush() else: self.stderr_logger.flush() if self.subprocess: if not self.stdout_logger: # Going to the console sys.stdout.flush() else: self.stdout_logger.flush() def __str__(self): # We only want to show the *different* environment variables def env_compare(env_orig, env_new): diff = {} for field, value in env_new.iteritems(): if not field in env_orig: diff[field] = value return diff env_diff = env_compare(os.environ, self.env) if self.env else {} if self.env_vars: for field, value in self.env_vars.iteritems(): env_diff[field] = value env_strs = [] for field, value in env_diff.iteritems(): env_strs.append("%s=%s" % (field, value)) cl = [] if env_strs: cl.append(" ".join(env_strs)) cl.append(self.executable) if self.arguments: cl.append(" ".join(self.arguments)) if self.subprocess: cl.append("(%s)" % self.subprocess.pid) return " ".join(cl) def __repr__(self): return self.__str__()