scram-package-monitor

#!/usr/bin/env python
import os, sys, json, socket, errno, re, urllib2, base64
from glob import glob
from datetime import datetime
from optparse import OptionParser
from os.path import basename, getctime, join
from os import getenv
from time import strftime, localtime, strptime, mktime
from hashlib import sha1
from commands import getstatusoutput

MAX_FILES_PUSH = 20
MAX_TIME_STOP = 120 #seconds
TIME_FORMAT = "%Y-%m-%dT%H:%M:%S.0"

def esReportPackages(results):
  # Silently exit if we cannot contact elasticsearch
  es_hostname = getenv("ES_HOSTNAME")
  es_auth = getenv("ES_AUTH")
  if not es_hostname and not es_auth:
    return

  url = "https://%s/_bulk" % (es_hostname)

  request = urllib2.Request(url)
  if es_auth:
    base64string = base64.encodestring(es_auth).replace('\n', '')
    request.add_header("Authorization", "Basic %s" % base64string)
  request.get_method = lambda: 'POST'
  data = "\n".join(results) + "\n"
  try:
    result = urllib2.urlopen(request, data=data)
  except urllib2.HTTPError, e:
    print e
    try:
      print result.read()
    except:
      pass

def create_dir(dir_to_create):
  try:
    os.makedirs(dir_to_create)
  except OSError as exception:
    if exception.errno != errno.EEXIST:
      raise exception

def check_pid(pid):
    """ Check For the existence of a unix pid. """
    try:
        os.kill(pid, 0)
    except OSError:
        return False
    else:
        return True

if __name__ == "__main__":
  WORKSPACE = os.getenv("WORKSPACE", "./")
  WORK_DIR = join(WORKSPACE, "pkg_mon")
  SOURCE_DIR = join(os.getenv("CMSSW_BASE", "./"), "src")

  parser = OptionParser(usage="%prog <-s|-e> -p <package>")
  parser.add_option("-s","--start", dest="start", action="store_true",
                    help="Building started for package", default=True)
  parser.add_option("-e","--stop", dest="start", action="store_false",
                    help="Building done for package", default=True)
  parser.add_option("-f","--force", dest="force", action="store_true",
                    help="Force pushing", default=False)
  parser.add_option("-p","--package", dest="pkg_name",
                    help="Package name to track", default=None)
  parser.add_option("-n", "--dry-run", dest="dryrun", action="store_true",
                    help="Do not push files to server", default=False)
  opts, args = parser.parse_args()

  CMSSW_VERSION = os.getenv("CMSSW_VERSION", "unknown")
  SCRAM_ARCH = os.getenv("SCRAM_ARCH", "unknown")
  INDEX_NAME = strftime("ib-scram-stats-%Y.%m.%d")
  defaults = { "hostname": socket.gethostname(),
               "scram_arch": SCRAM_ARCH,
               "cmssw_version": CMSSW_VERSION,
             }

  if (not opts.pkg_name) and opts.start:
    results = []
    for p in glob(join(SOURCE_DIR, "*/*")):
      pkg = re.match(SOURCE_DIR+"/(.+)", p).groups()[0]
      data = {"package": pkg}
      data.update(defaults)
      h = sha1(pkg + SCRAM_ARCH + CMSSW_VERSION).hexdigest()
      header = { "index" : { "_index" : INDEX_NAME,
                                     "_type" : "cmssw_pkg_times",
                                     "_id": h}
               }
      results += [json.dumps(header), json.dumps(data)]
    if opts.dryrun:
      print "Dry run specified, what I would have sent: \n" + "\n".join(results)
    else:
      esReportPackages(results)
    exit(0)

  try:
    thread_id = os.fork()
  except OSError, e:
    print "Error while forking"
    sys.exit(0)
  if not thread_id == 0:
    sys.exit(0)

  pkg_name = opts.pkg_name
  create_dir(WORK_DIR)

  # We look for stale "push" dirs, move all the files back in the workdir and
  # remove them.
  pushDirs = [join(WORKSPACE, d) for d in glob(join(WORKSPACE, "push_dir_*"))]
  pids = [re.sub(".*/push_dir_", "", p) for p in pushDirs]
  stalePids = [p for p in pids if not check_pid(int(p))]
  staleDirs = [join(WORKSPACE, "push_dir_" + p) for p in stalePids]
  for d in staleDirs:
    for f in glob(join(d, "*")):
      try:
        os.renames(f, join(WORK_DIR, basename(f)))
      except:
        pass
    try:
      os.rmdir(d)
    except:
      pass

  # Create the file for the current invocation.
  if pkg_name:
    prefix = opts.start and strftime("start_%s-") or strftime("stop_%s-")
    filename = prefix + pkg_name.replace("/",":")
    while(True):
      try:
        open(join(WORK_DIR, filename), "a").close()
        break
      except:
        create_dir(WORK_DIR)

  # When enough files are ready to push or enough time has passed, start
  # pushing to elasticsearch.
  try:
    fileLimitReached = len([f for f in os.listdir(WORK_DIR)]) > MAX_FILES_PUSH
  except OSError as exception:
    if exception.errno != errno.EEXIST:
      raise exception
    else:
      sys.exit(0)
  timestamps = sorted([int(basename(x).split("_")[1].split("-")[0])
                       for x in glob(join(WORK_DIR, "*"))])
  if len(timestamps) < 2:
    sys.exit(0)
  # Time passed between oldest file in the directory and the newest
  timeLimitReached = int(timestamps.pop()) - int(timestamps[0]) 
  limitsReached = fileLimitReached or timeLimitReached
  #Can force the push with force option or by calling end package without a
  # package
  if not (limitsReached or opts.force or not opts.pkg_name):
    sys.exit(0)

  push_dir = join(WORKSPACE, "push_dir_"+str(os.getpid()))
  try:
    os.rename(WORK_DIR, push_dir)
  except:
    sys.exit(0)

  results = []
  removables = []
  RE_FILE = "(start|stop)_([0-9]+)-(.*)"
  pushables = [f.replace(":", "/") for f in os.listdir(push_dir)]
  info = [re.match(RE_FILE, f).groups() for f in pushables]
  m = re.match("(.*)_(20[0-9]{2}-[0-9]{2}-[0-9]{2}-[0-9]{4})", CMSSW_VERSION)
  if m:
    defaults["cmssw_queue"] = m.group(1)
    defaults["@timestamp"] = strftime("%Y-%m-%dT%H:%M:00.0",
                                      strptime(m.group(2), "%Y-%m-%d-%H%M"))
  starts = dict([(x[2],int(x[1])) for x in info if x[0] == "start"])
  stops = dict([(x[2],int(x[1])) for x in info if x[0] == "stop"])
  packages = set(x[2] for x in info)
  results = []
  for x in packages:
    h = sha1(x + SCRAM_ARCH + CMSSW_VERSION).hexdigest()
    header = { "index" : { "_index" : INDEX_NAME,
                           "_type" : "cmssw_pkg_times",
                           "_id": h}
             }
    data = {"package": x}
    data.update(defaults)
    startTime = starts.get(x, None)
    stopTime = stops.get(x, None)
    if startTime:
      data["start"] = strftime(TIME_FORMAT,localtime(startTime))
    if stopTime:
      data["stop"] = strftime(TIME_FORMAT,localtime(stopTime))

    if startTime and stopTime:
      data["diff"] = stopTime - startTime
      removables.append(x.replace("/",":"))
    results += [json.dumps(header), json.dumps(data)]

  # Actually do the push to ES.
  if opts.dryrun:
    print "Dry run specified, what I would have sent:\n" + "\n".join(results)
  else:
    esReportPackages(results)

  for x in removables:
    cmd = "find %s -name \"*%s\" -delete" % (push_dir, x)
    err, out = getstatusoutput(cmd)
  for f in os.listdir(push_dir):
    while(True):
      try:
        os.renames(join(push_dir, f), join(WORK_DIR, f))
        break
      except:
        pass
  try:
    os.rmdir(push_dir)
  except:
    pass