path: root/seivot

#!/usr/bin/python
# Copyright 2011  Lars Wirzenius
# 
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import cliapp
import ConfigParser
import logging
import os
import shutil
import subprocess
import tempfile
import time


class Measurement(object):

    def __init__(self, time_output):
        fields = [float(x) for x in time_output.splitlines()]
        self.user = fields[0]
        self.system = fields[1]
        self.real = fields[2]
        self.maxrss = fields[3]
        self.new_data = 0


def runcmd(argv, **kwargs):
    logging.debug('run: %s %s' % (argv, kwargs))
    fd, timings = tempfile.mkstemp()
    time_argv = ['/usr/bin/time', 
                 '-o', timings, 
                 '--format', '%U\n%S\n%e\n%M']
    p = subprocess.Popen(time_argv + argv, **kwargs)
    out, err = p.communicate()
    os.remove(timings)
    data = os.read(fd, 1024**2)
    os.close(fd)
    if p.returncode != 0:
        raise cliapp.AppException('command failed: %s\n%s' % (argv, err))

    return Measurement(data), out


# Clear Linux kernel buffer and inode caches.
# See http://linux-mm.org/Drop_Caches for details.
def drop_caches():
    def sudo_tee(status):
        p = subprocess.Popen(['sudo', '-p', 'Password (for clearing cache): ',
                              'tee', '/proc/sys/vm/drop_caches'],
                             stdin=subprocess.PIPE, stdout=subprocess.PIPE)
        out, err = p.communicate('%s\n' % status)
        if p.returncode != 0:
            raise cliapp.AppException('failed to clear cache')

    logging.debug('clearing Linux kernel cache')
    sudo_tee(3)
    sudo_tee(0)


class BackupProgram(object):

    name = None
    
    def __init__(self, live_data, repo, settings):
        self.live_data = live_data
        self.repo = repo
        self.settings = settings

    def set_meta(self, cp):
        '''Set [meta] fields in report.
        
        These might depend on the program. For example, if running
        something from a version control branch, a subclass might
        record the revision id here.
        
        '''

    def prepare(self):
        '''Prepare program for benchmark.
        
        This might, for example, compile it.
        
        '''

    def backup(self, nth_gen):
        '''Run a backup from live_data to repo.
        
        This should start a new generation, in whatever way is 
        most appropriate for the backup program.
        
        '''
        
    def list_files(self, nth_gen):
        '''This should retrieve a list of all files in a generation.
        
        The list should be written to /dev/null.
        
        '''
        
    def restore(self, nth_gen, target_dir):
        '''Restore all files in a generation, to a target directory.'''
        
    def forget(self, nth_gen):
        '''Remove a given generation.'''


class Obnam(BackupProgram):

    name = 'obnam'
    
    @property
    def _cmd(self):
        if self.settings['obnam-branch']:
            return './obnam'
        else:
            return 'obnam'
    
    @property
    def _branch(self):
        return self.settings['obnam-branch'] or None

    @property
    def _revno(self):
        timings, revno = runcmd(['bzr', 'revno'], cwd=self._branch,
                                stdout=subprocess.PIPE)
        return revno.strip()
    
    @property
    def _larch_branch(self):
        return self.settings['larch-branch'] or None

    @property
    def _larch_revno(self):
        timings, revno = runcmd(['bzr', 'revno'], cwd=self._larch_branch,
                                stdout=subprocess.PIPE)
        return revno.strip()

    def _run(self, args, nth_gen, **kwargs):
        cmd = [self._cmd, 
               '--log', '/dev/null',
               '--repository', self.repo,
               '--weak-random']
        env = dict(os.environ)
        if self.settings['obnam-profile']:
            fd, env['OBNAM_PROFILE'] = tempfile.mkstemp()
            os.close(fd)
        if self._larch_branch:
            env['PYTHONPATH'] = self._larch_branch
        if self.settings['encrypt-with']:
            cmd += ['--encrypt-with', self.settings['encrypt-with']]
        result = runcmd(cmd + args, cwd=self._branch, env=env, **kwargs)
        if self.settings['obnam-profile']:
            namepattern = {
                'gen': str(nth_gen),
                'op': args[0],
            }
            for order in ['cumulative', 'time']:
                namepattern['order'] = order
                name = self.settings['obnam-profile'] % namepattern
                f = open(name, 'w')
                runcmd(['viewprof', env['OBNAM_PROFILE'], order], stdout=f)
                f.close()
            os.remove(env['OBNAM_PROFILE'])
        return result

    def prepare(self):
        if self._branch:
            logging.info('Building obnam in %s' % self._branch)
            runcmd(['make'], cwd=self._branch)
    
    def backup(self, nth_gen):
        return self._run(['backup', self.live_data], nth_gen)[0]
   
    def _genid(self, nth_gen):
        timings, out = self._run(['genids'], nth_gen, stdout=subprocess.PIPE)
        return out.splitlines()[nth_gen]
   
    def list_files(self, nth_gen):
        devnull = os.open('/dev/null', os.O_WRONLY)
        timings, out = self._run(['ls', self._genid(nth_gen)], nth_gen, 
                                 stdout=devnull)
        os.close(devnull)
        return timings

    def restore(self, nth_gen, target_dir):
        return self._run(['restore', '--to', target_dir, 
                           '--generation', self._genid(nth_gen)], nth_gen)[0]
    
    def forget(self, nth_gen, gen):
        return self._run(['forget', self._genid(gen)], nth_gen)[0]
    
    def set_meta(self, cp):
        if self._branch:
            cp.set('meta', 'revision', self._revno)
        if self._larch_branch:
            cp.set('meta', 'larch-revision', self._larch_revno)


class BackupProgramFactory(object):

    programs = [Obnam]
    
    def names(self):
        return [p.name for p in self.programs]

    def new(self, name, **kwargs):
        for p in self.programs:
            if p.name == name:
                return p(**kwargs)


class Report(object):

    def __init__(self, program):
        self.program = program
        self.measurements = dict()
        
    @property
    def generations(self):
        gens = set()
        for op in self.measurements:
            for gen in self.measurements[op]:
                gens.add(gen)
        return gens

    @property
    def operations(self):
        return self.measurements.keys()
        
    def add_measurement(self, op, gen, measurement):
        if op not in self.measurements:
            self.measurements[op] = dict()
        self.measurements[op][gen] = measurement

    def get_measurement(self, op, gen):
        return self.measurements[op][gen]
        
    def format(self, fp):
        cp = ConfigParser.ConfigParser()

        cp.add_section('meta')
        cp.set('meta', 'program', self.program.name)
        self.program.set_meta(cp)
        
        for gen in self.generations:
            section = str(gen)
            cp.add_section(section)
            for op in self.operations:
                m = self.get_measurement(op, gen)
                for field in ['user', 'system', 'real', 'maxrss']:
                    cp.set(section, '%s.%s' % (op, field), 
                           '%.1f' % getattr(m, field))
                cp.set(section, '%s.new-data' % op, m.new_data)
        
        cp.write(fp)


class Seivot(cliapp.Application):

    def add_settings(self):
        self.factory = BackupProgramFactory()

        self.settings.add_choice_setting(['program'], self.factory.names(),
                                         'program to benchmark (%default)')

        self.settings.add_integer_setting(['generations'],
                                          'total number of generations to '
                                            'measure (%default)',
                                          metavar='COUNT',
                                          default=5)
        self.settings.add_bytesize_setting(['initial-data'],
                                           'size of initial live data '
                                                '(%default)',
                                           metavar='SIZE',
                                           default=1024)
        self.settings.add_bytesize_setting(['incremental-data'],
                                           'add SIZE live data for '
                                                'additional generations '
                                                '(%default)',
                                           metavar='SIZE',
                                           default=1024)

        self.settings.add_string_setting(['obnam-branch'],
                                         'bzr branch from which to run obnam '
                                            '(default is installed obnam)')
        self.settings.add_string_setting(['larch-branch'],
                                         'bzr branch from which to use larch '
                                            '(default is installed larch)')
        self.settings.add_string_setting(['obnam-profile'],
                                          'store Python profiling output '
                                            'in files named after NAMEPATTERN '
                                            '(no profiling, unless set); '
                                            '%(foo)s in pattern gets filled '
                                            'in, where foo is op (for '
                                            'backup/restore/etc), gen, or '
                                            'order (cumulative/time)',
                                          metavar='NAMEPATTERN',
                                          default='')
        self.settings.add_string_setting(['encrypt-with'],
                                         'encrypt backups with KEYID',
                                         metavar='KEYID')
                                         
        self.settings.add_boolean_setting(['drop-caches'],
                                         'clear Linux kernel cache before '
                                           'running commands (will ask for '
                                           'sudo pasword')

    def process_args(self, args):
        progname = self.settings['program']
        logging.info('Benchmarking: %s' % progname)
        
        generations = self.settings['generations']

        self.tempdir = tempfile.mkdtemp()
        logging.info('tempdir: %s' % self.tempdir)
        self.live_data = os.path.join(self.tempdir, 'data')
        self.repo = os.path.join(self.tempdir, 'repo')

        prog = self.factory.new(progname, live_data=self.live_data,
                                repo=self.repo,
                                settings=self.settings)
        prog.prepare()
        
        self.report = Report(prog)

        self.generate_live_data(self.live_data, self.settings['initial-data'])
        self.measure(prog.backup, 0, self.settings['initial-data'])

        for i in range(1, generations):
            self.generate_live_data(self.live_data,
                                    self.settings['incremental-data'])
            self.measure(prog.backup, i, self.settings['incremental-data'])

        for i in range(generations):
            self.measure(prog.list_files, i, 0)

        for i in range(generations):
            target_dir = os.path.join(self.tempdir, 'restored')
            os.mkdir(target_dir)
            self.measure(prog.restore, i, 0, target_dir=target_dir)
            shutil.rmtree(target_dir)

        for i in range(generations):
            # Since we remove oldest first, we always remove the 0th
            # generation, not the ith one.
            self.measure(prog.forget, i, 0, gen=0)

        self.cleanup()
        
        self.report.format(self.output)

    def generate_live_data(self, where, size):
        logging.info('Generating %d bytes live data' % size)
        runcmd(['genbackupdata', where, '--create', str(size)])

    def measure(self, func, nth_gen, new_data, **kwargs):
        logging.info('Measuring %s gen %d' % (func.__name__, nth_gen))
        drop_caches()
        measurement = func(nth_gen, **kwargs)
        measurement.new_data = new_data
        self.report.add_measurement(func.__name__, nth_gen, measurement)

    def cleanup(self):
        logging.info('Removing temporary directory %s' % self.tempdir)
        shutil.rmtree(self.tempdir)


if __name__ == '__main__':
    Seivot().run()