benchmark_pijul.py

#!/usr/bin/env python3
"""Benchmark Pijul's CLI on the shared deterministic single-file histories."""
import argparse
import atexit
import hashlib
import json
import os
import platform
import shutil
import statistics
import subprocess
import tarfile
import tempfile
import time
from pathlib import Path

from benchmark_fixtures import versions

ROOT = Path(__file__).resolve().parent


def digest(data):
    return hashlib.sha256(data).hexdigest()


def summarize(samples):
    ordered = sorted(samples)
    return {'median_ms': statistics.median(samples), 'min_ms': min(samples),
            'max_ms': max(samples), 'p95_ms': ordered[min(len(ordered) - 1, int(len(ordered) * .95))],
            'samples_ms': samples}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--pijul', default=shutil.which('pijul'))
    parser.add_argument('--seed', type=int, default=17879179514528422438)
    parser.add_argument('--versions', type=int, default=100)
    parser.add_argument('--runs', type=int, default=3)
    parser.add_argument('--reads', type=int, default=3)
    args = parser.parse_args()
    if not args.pijul or not Path(args.pijul).is_file():
        parser.error('provide an existing Pijul executable with --pijul')
    if min(args.versions, args.runs, args.reads) < 1:
        parser.error('versions, runs, and reads must be positive')
    binary = str(Path(args.pijul).resolve())
    directory = Path(tempfile.mkdtemp(prefix='pijul-benchmark-', dir=ROOT))
    home = directory / 'home'
    home.mkdir()
    env = os.environ.copy()
    env.update(HOME=str(home), XDG_CONFIG_HOME=str(home / 'config'),
               XDG_CACHE_HOME=str(home / 'cache'), XDG_DATA_HOME=str(home / 'data'),
               LC_ALL='C', TZ='UTC', USER='benchmark')
    env.pop('PIJUL_REPOSITORY', None)
    env.pop('SSH_AUTH_SOCK', None)
    key = home / 'benchmark-key'
    ssh_tools = {name: str(Path(shutil.which(name)).resolve()) for name in ('ssh-keygen', 'ssh-agent', 'ssh-add')}
    subprocess.run([ssh_tools['ssh-keygen'], '-q', '-t', 'ed25519', '-N', '', '-f', str(key)],
                   check=True, env=env, stdin=subprocess.DEVNULL)
    agent_output = subprocess.check_output([ssh_tools['ssh-agent'], '-s'], env=env, text=True)
    for part in agent_output.split(';'):
        part = part.strip()
        if part.startswith(('SSH_AUTH_SOCK=', 'SSH_AGENT_PID=')):
            name, value = part.split('=', 1)
            env[name] = value
    atexit.register(lambda: subprocess.run([ssh_tools['ssh-agent'], '-k'], env=env,
                                           stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL))
    subprocess.run([ssh_tools['ssh-add'], str(key)], check=True, env=env,
                   stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE)

    def call(folder, *command):
        start = time.perf_counter_ns()
        result = subprocess.run([binary, *map(str, command)], cwd=folder, env=env,
                                stdin=subprocess.DEVNULL, stdout=subprocess.PIPE,
                                stderr=subprocess.PIPE, timeout=300)
        ms = (time.perf_counter_ns() - start) / 1e6
        if result.returncode:
            raise RuntimeError(f'{command!r} failed ({result.returncode}): {result.stderr.decode(errors="replace")}\n{result.stdout.decode(errors="replace")}')
        return ms, result.stdout

    _, version = call(directory, '--version')
    report = {'completed': False, 'seed': args.seed, 'platform': platform.platform(),
              'python': platform.python_version(), 'arguments': vars(args),
              'script_sha256': digest(Path(__file__).read_bytes()),
              'fixture_generator_sha256': digest((ROOT / 'benchmark_fixtures.py').read_bytes()),
              'binary': {'path': binary, 'sha256': digest(Path(binary).read_bytes()),
                         'version': version.decode().strip()},
              'ssh_tools': {name: {'path': path, 'sha256': digest(Path(path).read_bytes())} for name, path in ssh_tools.items()},
              'directory': str(directory), 'cases': [],
              'protocol': ('Fresh isolated repository per trial. Each working file write is excluded; timed writes include Pijul process startup, record, and initial add. '
                           'Latest read timing includes Pijul startup, reconstructing the entire repository into a gzip-compressed tar archive, and writing it to disk; archive extraction and hash verification are excluded. '
                           'Historical versions are reconstructed using archive --state and verified outside timing. Repository size is logical bytes in .pijul only, excluding the working copy. '
                           'Initialization, isolated SSH agent and ephemeral signing-key setup, fixture generation, log collection, and verification are excluded. Warm caches and background load are uncontrolled.'),
              'limitations': 'Pijul archive writes a compressed whole-repository tar file, unlike tools returning raw file contents on stdout; read times are not directly equivalent. CLI startup, storage, metadata and durability policies differ.'}
    result_path = directory / 'results.json'

    def save():
        result_path.write_text(json.dumps(report, indent=2) + '\n')

    save()
    for size in (4096, 131072):
        for workload in ('binary', 'text'):
            for trial in range(args.runs):
                folder = directory / f'{workload}-{size}-{trial}'
                folder.mkdir()
                call(folder, 'init')
                case = {'workload': workload, 'initial_bytes': size, 'trial': trial,
                        'fixture_sha256': [], 'tools': {'pijul': {'write_ms': [], 'checkpoints': []}}}
                report['cases'].append(case)
                tool = case['tools']['pijul']
                for index, data in enumerate(versions(args.seed, size, workload, args.versions)):
                    h = digest(data)
                    if h in case['fixture_sha256']:
                        raise RuntimeError('fixture contains a duplicate version')
                    case['fixture_sha256'].append(h)
                    (folder / 'content').write_bytes(data)
                    elapsed = 0
                    if index == 0:
                        elapsed, _ = call(folder, 'add', 'content')
                    ms, _ = call(folder, 'record', '-a', '-m', f'version {index}', '--author', 'benchmark')
                    tool['write_ms'].append(elapsed + ms)
                _, log = call(folder, 'log', '--state', '--output-format', 'json')
                entries = json.loads(log)
                # Messages map entries to fixture indices; do not assume log order.
                states = {}
                for entry in entries:
                    message = entry.get('message', '')
                    if message.startswith('version '):
                        states[int(message.split()[1])] = entry['state']
                if set(states) != set(range(args.versions)):
                    raise RuntimeError(f'log does not contain exactly {args.versions} benchmark versions')
                case['revision_states'] = [states[i] for i in range(args.versions)]
                archive = directory / 'read.tar.gz'

                def retrieve(state=None):
                    command = ['archive', '-o', archive]
                    if state is not None:
                        command += ['--state', state]
                    ms, _ = call(folder, *command)
                    with tarfile.open(archive, 'r:gz') as tar:
                        f = tar.extractfile('content')
                        if f is None:
                            raise RuntimeError('content missing from archive')
                        output = f.read()
                    archive.unlink()
                    return ms, digest(output)

                reads = []
                for _ in range(args.reads):
                    ms, actual = retrieve()
                    if actual != case['fixture_sha256'][-1]:
                        raise RuntimeError('incorrect latest reconstruction')
                    reads.append(ms)
                repo_size = sum(p.stat().st_size for p in (folder / '.pijul').rglob('*') if p.is_file())
                tool['checkpoints'].append({'versions': args.versions, 'repository_logical_bytes': repo_size,
                    'write_last_100': summarize(tool['write_ms'][-100:]),
                    'cumulative_write_ms': sum(tool['write_ms']), 'reads': {'latest': summarize(reads)}})
                for index in range(args.versions):
                    _, actual = retrieve(states[index])
                    if actual != case['fixture_sha256'][index]:
                        raise RuntimeError(f'incorrect reconstruction of revision {index}')
                tool['verified_revisions'] = args.versions
                save()
                print(f'{workload}, {size} bytes, trial {trial + 1}, {args.versions} versions complete', flush=True)
    report['completed'] = True
    save()
    print(f'Results: {result_path}', flush=True)


if __name__ == '__main__':
    main()