buildenv

#!/usr/bin/env python3

import collections
import datetime
import json
import os.path
import pickle
import re
import sh
import shutil
import sys

SELF = sys.argv[0]
ROOT = os.path.dirname(SELF)

def json_converter(o):
  if isinstance(o, datetime.datetime):
    return o.isoformat()
  raise TypeError('%s is not JSON serializable' % o)

def mv(f, t):
  os.rename(f, t)

def mkpath(p):
  try:
    os.makedirs(p)
  except OSError as e:
    if e.args[0] != 17: # EEXIST
      raise

def copy(f, t):
  mkpath(os.path.dirname(t))
  sh.cp('-a', f, t)

def rm(p):
  try:
    if not os.path.isdir(p) or os.path.islink(p):
      os.unlink(p)
    else:
      shutil.rmtree(p)
  except OSError as e:
    if e.args[0] != 2: # ENOENT
      raise

class CLError(Exception):
  pass

class State:

  def __init__(self):
    self.__parser = self.__options_parser()
    self.__opts = self.__parser.parse_args()
    self.__branch = self.__opts.branch
    self.__cache = None
    self.__git = None

  @property
  def state_path(self):
    return '%s/.state' % self.__opts.build_trees_root

  def __enter__(self):
    mkpath(self.__opts.build_trees_root)
    try:
      with open(self.state_path, 'rb') as f:
        self.__state = pickle.load(f)
      if self.__state is None:
        self.__state = {}
    except OSError:
      self.__state = {}
    self.__validate_state()
    return self

  def __exit__(self, t, v, tb):
    if not self.__opts.dry_run and t is None:
      tmp = '%s.tmp' % self.state_path
      try:
        with open(tmp, 'wb') as f:
          pickle.dump(self.__state, f)
      except:
        rm(tmp)
        raise
      else:
        mv(tmp, self.state_path)

  def __validate_state(self):
    trees = self.__state.setdefault('build-trees', {})
    for b in list(trees):
      tree = trees[b]
      if not os.path.exists(tree['path']):
        print('Missing build tree %s, evicting from state' % tree['path'])
        del trees[b]

  def __options_parser(self):
    import argparse
    # Parse options
    parser = argparse.ArgumentParser(add_help = False,
                                     description = 'To infinity!')
    opt = lambda *a, **kw: parser.add_argument(*a, **kw)
    opt('-h', '--help', action = 'help', default = argparse.SUPPRESS,
        help = 'Show this help message and exit',)
    opt('-b', '--branch', type = str, help = 'Built branch',)
    opt('-B', '--build-trees-root', type = str, default = '%s/_builds' % ROOT,
        help = 'Build trees root directory',)
    opt('-p', '--pick-build-tree', action = 'store_true',
        help = 'Pick adequate build tree',)
    opt('-d', '--cache-download-build-tree', action = 'store_true',
        help = 'Download build tree cache',)
    opt('-u', '--cache-upload-build-tree', action = 'store_true',
        help = 'Upload build tree cache',)
    opt('-c', '--cleanup-build-trees', action = 'store_true',
        help = 'Garbage collect unused build trees',)
    opt('-C', '--cleanup-threshold', type = int, default = 10,
        help = 'Minimum headroom (in GB) below which build trees are cleaned',)
    opt('-n', '--cache-namespace', type = str,
        help = 'Subdirectory where to locate cache',)
    opt('-i', '--cache-id', type = str,
        help = 'Concurrent cache unique id',)
    opt('-l', '--link-build-tree', type = str,
        help = 'Create symlink to picked build tree',)
    opt('-s', '--show-build-tree', action = 'store_true',
        help = 'Pick adequate build tree',)
    opt('-S', '--source-tree', type = str, default = ROOT,
        help = 'Source tree location',)
    opt('-D', '--dry-run', action = 'store_true',
        help = 'Do not modify anything',)
    opt('-v', '--verbose', action = 'store_true',
         help = 'Increase verbosity',)
    opt('-e', '--exclude', type = str,
        help = 'Exclude given pattern from rsync',)
    return parser

  @property
  def parser(self):
    return self.__parser

  @property
  def source_tree(self):
    return self.__opts.source_tree

  @property
  def git(self):
    if self.__git is None:
      self.__git = sh.git.bake(_cwd = self.source_tree,
                               _tty_in = False,
                               _tty_out = False)
    return self.__git

  @property
  def branch(self):
    if self.__branch is None:
      self.__branch = self.git(
        'rev-parse', '--abbrev-ref', 'HEAD').strip()
    return self.__branch

  def run(self):
    if self.__opts.pick_build_tree:
      tree = self.pick_build_tree()
      symlink = self.__opts.link_build_tree
      if symlink is not None:
        rm(symlink)
        os.symlink(tree['path'], symlink)
      return tree
    elif self.__opts.show_build_tree:
      return self.show_build_tree()
    elif self.__opts.cache_download_build_tree:
      return self.download_build_tree()
    elif self.__opts.cache_upload_build_tree:
      return self.upload_build_tree()
    elif self.__opts.cleanup_build_trees:
      return self.garbage_collect_build_trees()
    else:
      raise CLError('specify an action')

  @property
  def build_tree(self):
    return '%s/%s' % (self.__opts.build_trees_root, self.branch)

  def pick_build_tree(self):
    state = self.__state
    build_trees = state.setdefault('build-trees', {})
    branch = self.branch
    current_rev = self.git('rev-parse', 'HEAD').strip()
    previous = build_trees.get(branch)
    if previous is not None:
      if os.path.exists(previous['path']):
        previous_rev = previous['rev']
        previous['rev'] = current_rev
        previous['last_use'] = datetime.datetime.utcnow()
        return {
           'status': 'existing build tree',
           'path': previous['path'],
           'rev': previous_rev,
         }
      else:
        del build_trees[branch]
        previous = None
    clone_d = None
    clone = None
    clone_branch = None
    for b, tree in build_trees.items():
      if not os.path.exists(tree['path']):
        continue
      rev = tree['rev']
      base = self.git('merge-base', rev, current_rev).strip()
      def distance(r):
        return int(sh.wc(self.git('log', '%s..%s' % (base, r),
                                  '--pretty=oneline'), '-l'))
      d = distance(current_rev) + distance(rev)
      if clone is None or d < clone_d:
        clone = tree
        clone_branch = b
        clone_d = d
    build_trees[branch] = {
      'creation': datetime.datetime.utcnow(),
      'last_use': datetime.datetime.utcnow(),
      'path': self.build_tree,
      'rev': current_rev,
    }
    if clone is not None:
      if not self.__opts.dry_run:
        tmp = '%s.tmp' % self.build_tree
        try:
          copy(tree['path'], tmp)
          fr = os.path.realpath(tree['path'])
          fr = '%s/.drake/%s' % (tmp, fr)
          to = os.path.realpath(self.build_tree)
          to = '%s/.drake/%s' % (tmp, to)
          mkpath(os.path.dirname(to))
          rm(to)
          if os.path.exists(fr):
            mv(fr, to)
          else:
            print('%s: missing drake directory: %s' % (sys.argv[0], fr),
                  file = sys.stderr)
        except:
          rm(tmp)
          raise
        else:
          mv(tmp, self.build_tree)
      return {
        'status': 'cloned from %s' % clone_branch,
        'path': self.build_tree,
        'rev': tree['rev'],
      }
    else:
      if not self.__opts.dry_run:
        mkpath(self.build_tree)
      return {
        'status': 'new build tree',
        'path': self.build_tree,
        'rev': None,
      }

  def show_build_tree(self):
    build_trees = self.__state.setdefault('build-trees', {})
    previous = build_trees.get(self.branch)
    if previous is not None:
      return previous
    else:
      return None

  def disk_space(self):
    '''Available disk space in GB.'''
    statvfs = os.statvfs('.')
    res = int(statvfs.f_frsize * statvfs.f_bavail / 1024 ** 3)
    return res

  def garbage_collect_build_trees(self):
    '''Remove old build trees if we lack space.  Remove per day, starting
    from 10d olds, down to 2d, stopping as soon as we have the desired
    headroom.

    Stopping at two days, not zero, means that we may have builds that
    will fail because we lack space.  But we prefer this, as it sends
    a signal for someone to pay attention to what is going on.  Had we
    proceeded to 0d, we may not notice that the caching of build is
    completely disabled.'''

    cleaned = []
    state = self.__state
    trees = state.setdefault('build-trees', {})
    def clean_older(delay):
      for b in list(trees):
        tree = trees[b]
        if datetime.datetime.utcnow() - tree['last_use'] > delay:
          cleaned.append((b, tree['last_use']))
          del trees[b]
          try:
            shutil.rmtree(tree['path'])
          except FileNotFoundError:
            pass
    max_age = datetime.timedelta(days = 7)
    clean_older(max_age)
    while self.disk_space() < self.__opts.cleanup_threshold and \
          max_age >= datetime.timedelta(days = 3):
      max_age -= datetime.timedelta(days = 1)
      clean_older(max_age)
    return {
      'cleaned': cleaned,
    }

  cache_url = 'buildslave@cache.buildslave.infinit.sh'

  def cache(self):
    if self.__cache is None:
      self.__cache = sh.ssh.bake(self.cache_url,
                                 _tty_in = False,
                                 _tty_out = False)
    return self.__cache

  class Lock:

    def __init__(self, path):
      self.__path = '%s/rsync' % path

    def __enter__(self):
      sh.ssh(State.cache_url, 'lockfile-create', self.__path)

    def __exit__(self, t, v, tb):
      sh.ssh(State.cache_url, 'lockfile-remove', self.__path)

  @property
  def old_cache_dir(self):
    import getpass
    ns = self.__opts.cache_namespace or getpass.getuser()
    return 'cache/%s/%s' % (ns, self.branch)

  @property
  def cache_dir(self):
    import getpass
    ns = self.__opts.cache_namespace or getpass.getuser()
    return 'cache/%s/%s' % (ns, self.branch.replace('/', '_'))

  def rsync(self, s, d, exclude = None):
    v = self.__opts.verbose
    args = [
      '--archive',
      '--checksum',
      '--human-readable',
      # '--progress',
      '--recursive',
      '--update',
      '--delete',
    ]
    if exclude is not None:
      if isinstance(exclude, str):
        args += ['--exclude', exclude]
      else:
        for e in exclude:
          args += ['--exclude', e]
    args += [
      s, d,
    ]
    if v:
      args.append('--stats')
    res = sh.rsync(*args)
    if v:
      print(res)

  def download_build_tree(self):
    try:
      sh.ssh(self.cache_url, 'test', '-e', self.cache_dir)
    except sh.ErrorReturnCode:
      return {
        'status': 'no cached build tree',
      }
    else:
      if not self.__opts.dry_run:
        with State.Lock(self.cache_dir):
          self.rsync(
            '%s:%s/' % (self.cache_url, self.cache_dir),
            self.build_tree,
            exclude = self.__opts.exclude and self.__opts.exclude.split(',') or None)

    return {
      'status': 'downloaded build tree',
      'source': self.cache_dir,
      'destination': self.build_tree,
    }

  def upload_build_tree(self):
    d = self.cache_dir
    concurrent = self.__opts.cache_id is not None
    if concurrent:
      d = '%s-%s' % (d, self.__opts.cache_id)
    if not self.__opts.dry_run:
      sh.ssh(self.cache_url, 'mkdir', '-p', d)
      with State.Lock(d):
        self.rsync(
          '%s/' % self.build_tree,
          '%s:%s' % (self.cache_url, d),
           exclude = self.__opts.exclude and self.__opts.exclude.split(',') or None)
      if concurrent:
        sh.ssh(self.cache_url,
               'ln', '-sfn', os.path.basename(d), self.cache_dir)
    return {
      'status': 'uploaded build tree',
      'source': self.build_tree,
      'destination': d,
    }

try:
  with State() as state:
    res = state.run()
    if isinstance(res, dict):
      res = collections.OrderedDict(res)
    json.dump(res, sys.stdout, default = json_converter)
    print()
except CLError as e:
  state.parser.print_usage(file = sys.stderr)
  print('%s: command line error: %s' % (sys.argv[0], e),
        file = sys.stderr)
  exit(1)
except Exception as e:
  print('%s: fatal error: %s' % (sys.argv[0], e), file = sys.stderr)
  raise
  exit(1)