#!/usr/bin/env python
#
# This is a Condor job wrapper for dvdt. It is used to insert the dvdt monitoring
# tools, kickstart and resource_monitor, as wrappers for Condor jobs.
#
# To configure Condor to use this wrapper, modify the Condor configuration to
# include:
#
#     USER_JOB_WRAPPER = /path/to/dvdt_wrapper
#
# The configuration for this wrapper is below. At a minimum, you must specify
# the paths to kickstart and resource_monitor, and the data_dir where log files
# will be stored. You can also restrict the users whose jobs will be wrapped,
# and limit the percentage of jobs that will be wrapped.
#

# This is a list of the users whose jobs can be wrapped.
allow_owners = []

# This is a list of the users whose jobs cannot be wrapped.
deny_owners = []

# This is a list of users whose jobs cannot be wrapped with kickstart/ptrace
ptrace_owners = []

# This is the fraction of matching jobs that will be wrapped.
# 1.0 = all matching jobs
wrap_fraction = 1.0

# This is the monitor to use. The value should be either
# 'kickstart' or 'resource_monitor'.
monitor = "kickstart"

# This is the path to kickstart
kickstart = "/usr/bin/pegasus-kickstart"

# This is the path to resource_monitor
resource_monitor = "/Users/juve/Workspace/cctools/resource_monitor/src/resource_monitorv"

# This is the directory where output data is stored. Under this directory
# the script will create a directory for each schedd, and under that
# one directory for every group of 1000 jobs.
data_dir = "/submit/testjobs"

import os
import sys
import subprocess
import random
import platform
import syslog
import traceback
import errno
import time

class SysLogger(object):
    def __init__(self):
        syslog.openlog(os.path.basename(sys.argv[0]),
                       syslog.LOG_PID,
                       syslog.LOG_USER)

    def debug(self, message):
        syslog.syslog(syslog.LOG_DEBUG, message)

    def info(self, message):
        syslog.syslog(syslog.LOG_INFO, message)

    def notice(self, message):
        syslog.syslog(syslog.LOG_NOTICE, message)

    def warning(self, message):
        syslog.syslog(syslog.LOG_WARNING, message)

    def error(self, message):
        syslog.syslog(syslog.LOG_ERR, message)

    def critical(self, message):
        syslog.syslog(syslog.LOG_CRIT, message)

    def alert(self, message):
        syslog.syslog(syslog.LOG_ALERT, message)

    def emergency(self, message):
        syslog.syslog(syslog.LOG_EMERG, message)

    def exception(self, message=None):
        exc_info = sys.exc_info()

        lines = []

        # Add the message, if any
        if message is not None:
            message = str(message)
            if not message.endswith("\n"):
                message += "\n"
            lines.append(message)

        # Now add the exception and traceback
        if exc_info != (None, None, None):
            trace = traceback.format_exception(*exc_info)
            first = trace.pop()
            lines.append(first)
            lines.extend(trace)

        text = "".join(lines)
        syslog.syslog(syslog.LOG_ERR, text)

logger = SysLogger()

def getenv(name):
    value = os.getenv(name)
    if value is None:
        sys.stderr.write("%s is not defined in environment\n" % name)
        exit(1)
    return value

condor_wrapper_error_file = getenv("_CONDOR_WRAPPER_ERROR_FILE")
condor_job_ad = getenv("_CONDOR_JOB_AD")
condor_machine_ad = getenv("_CONDOR_MACHINE_AD")

def parse_classad(filename):
    classad = {}

    f = open(filename, "r")

    for l in f:
        if "=" not in l:
            continue

        l = l.rstrip()
        key, value = l.split(" = ", 1)

        # Strip quotes off of strings
        if value.startswith('"') and value.endswith('"'):
            value = value.strip('"')

        classad[key] = value

    f.close()

    return classad

def restart_job(message):
    # If an error occurs, then write to the error file
    # and the job will be retried

    logger.notice("Restarting job: %s" % message)

    f = open(condor_wrapper_error_file, "w")
    try:
        f.write(message)
    finally:
        f.close()

    exit(1)

class Wrapper(object):
    def __init__(self):
        self.command = sys.argv[1:]

        self.jobad = parse_classad(condor_job_ad)

        # Compute the classads we need
        self.GlobalJobId = self.jobad.get("GlobalJobId", "Unknown")
        self.ClusterId = self.jobad.get("ClusterId", "0")
        self.ProcId = self.jobad.get("ProcId", "0")
        self.NumShadowExceptions = self.jobad.get("NumShadowExceptions", "0")
        self.JobUniverse = self.jobad.get("JobUniverse", None)
        self.Owner = self.jobad.get("Owner", "")
        self.User = self.jobad.get("User", "nobody@localhost")
        self.Schedd = self.User.split("@")[1]
        self.JobId = "%s.%s" % (self.ClusterId, self.ProcId)

        # We group the logs into directories of 1000 jobs by schedd
        dirnum = (int(self.ClusterId) / 1000)
        dirname = "%d000s" % dirnum
        self.log_dir = os.path.join(data_dir, self.Schedd, dirname)

    def consider_job(self):
        logger.notice("Considering job %s" % self.GlobalJobId)

        if w.should_wrap_job():
            w.wrap_job()
        else:
            w.run_job()

    def should_wrap_job(self):

        # There can be no prior shadow exceptions for the job
        if self.NumShadowExceptions != "0":
            logger.notice("Not wrapping: NumShadowExceptions != 0: %s" %
                    self.NumShadowExceptions)
            return False

        # Has to be a vanilla universe job (5)
        if self.JobUniverse != "5":
            logger.notice("Not wrapping: JobUniverse != 5: %s" %
                    self.JobUniverse)
            return False

        # Only wrap users in the allow list
        if allow_owners:
            if self.Owner not in allow_owners:
                logger.notice("Not wrapping: Owner '%s' not in allow_owners" %
                        self.Owner)
                return False

        # Do not wrap users in the deny list
        if deny_owners:
            if self.Owner in deny_owners:
                logger.notice("Not wrapping: Owner '%s' in deny_owners" %
                        self.Owner)
                return False

        # Only wrap a fraction of matching jobs
        if wrap_fraction < 1.0:
            rnd = random.random()
            if wrap_fraction <= rnd:
                logger.notice("Not wrapping: Probability filter %f" %
                        wrap_fraction)
                return False

        # Wrap everything by default
        return True

    def create_log_dir(self):
        "Create the log dir if it doesn't exist"

        if os.path.isdir(self.log_dir):
            return

        logger.notice("Creating log_dir '%s'" % self.log_dir)

        # We have to set the umask to zero because makedirs applies it
        umask = os.umask(0)

        # We have to do this in a while loop in case of race conditions
        while not os.path.isdir(self.log_dir):
            try:
                # We allow anyone to access the directories and set
                # the sticky bit because we expect that multiple different
                # users will access the directory
                os.makedirs(self.log_dir, 01777)
            except OSError, e:
                # In case of race conditions
                if e.errno == errno.EEXIST:
                    logger.warning("Directory exists: %s" % self.log_dir)
                else:
                    raise e
            finally:
                # Restore the umask
                os.umask(umask)

    def wrap_job(self):
        "Run the job with one of the monitors"
        global monitor

        self.create_log_dir()

        if self.Owner in ptrace_owners:
            logger.warning("Owner is in ptrace_owners; using resource_monitor")
            monitor = 'resource_monitor'

        if platform.system() != "Linux":
            logger.warning("Platform is not Linux; using resource_monitor")
            monitor = 'resource_monitor'

        start = time.time()

        if monitor == 'kickstart':
            rc = self.wrap_job_kickstart()
        elif monitor == 'resource_monitor':
            rc = self.wrap_job_resource_monitor()
        else:
            logger.warning("Invalid monitor: %s" % monitor)
            self.run_job() # This should never return
            exit(1)

        finish = time.time()

        # Retry failed jobs that run for less than 1 second
        if (finish - start) < 1.0 and rc != 0:
            restart_job("Wrapped process failed after less than 1 second")

        exit(rc)

    def wrap_job_kickstart(self):
        logger.notice("Wrapping job with kickstart")

        infile = self.jobad.get("In", None)
        outfile = self.jobad.get("Out", None)
        errfile = self.jobad.get("Err", None)

        if infile is None:
            infile = "/dev/null"
        if outfile is None:
            outfile = "/dev/null"
        if errfile is None:
            errfile = "/dev/null"

        # Path to invocation record
        inv = os.path.join(self.log_dir, "kickstart-"+self.JobId+".xml")

        # Construct a new command line with kickstart
        newcommand = [kickstart, "-l", inv, "-i", infile, "-o", outfile, "-e", errfile]
        if platform.system() == "Linux":
            newcommand.append("-t")
        newcommand.extend(self.command)

        logger.info("Running %s" % " ".join(newcommand))

        proc = subprocess.Popen(newcommand, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

        stdout, stderr = proc.communicate()

        if len(stdout) > 0:
            logger.error("Kickstart produced stdout:\n%s" % stdout)
        if len(stderr) > 0:
            logger.error("Kickstart produced stderr:\n%s" % stderr)
        if len(stdout) > 0 or len(stderr) > 0:
            restart_job("Kickstart produced stdio")

        return proc.returncode

    def wrap_job_resource_monitor(self):
        logger.notice("Wrapping job with resource_monitor")

        # Resource monitor file name template
        template = os.path.join(self.log_dir, "resource_monitor-"+self.JobId)

        # Construct a new command line with resource_monitor
        newcommand = [resource_monitor, "-O", template, "--"]
        newcommand.extend(self.command)

        logger.info("Running %s" % " ".join(newcommand))

        proc = subprocess.Popen(newcommand)

        proc.wait()

        return proc.returncode

    def run_job(self):
        logger.notice("Not wrapping job")

        try:
            exe = self.command[0]
            os.execv(exe, self.command)
            sys.stderr.write("execv: %s: unable to run program\n", exe)
            exit(1)
        except OSError, e:
            sys.stderr.write("execv: %s: %s\n" % (exe, e))
            exit(1)

if __name__ == '__main__':
    if len(sys.argv) < 2:
        sys.stderr.write("Usage: %s COMMAND [ARG...]\n" % sys.argv[0])
        exit(1)

    try:
        w = Wrapper()
        w.consider_job()
    except Exception, e:
        logger.exception("Error wrapping job")
        restart_job(str(e))
        exit(1)

