bench_pack_recover.py

#!/usr/bin/env python3
"""Compare pack and recover on identical records with two NAB binaries."""
import argparse
import hashlib
import json
import platform
import statistics
import subprocess
import tempfile
import time
from pathlib import Path

parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('archive', type=Path)
parser.add_argument('old', type=Path)
parser.add_argument('new', type=Path)
parser.add_argument('output', type=Path)
parser.add_argument('--runs', type=int, default=5)
a = parser.parse_args()
old, new = a.old.resolve(strict=True), a.new.resolve(strict=True)
archive = a.archive.resolve(strict=True)

def sha(path):
    return hashlib.sha256(path.read_bytes()).hexdigest()

def command(args):
    start = time.perf_counter_ns()
    r = subprocess.run([str(x) for x in args], stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)
    return (time.perf_counter_ns() - start) / 1e6, r.stdout

log = subprocess.check_output([str(new), 'log', str(archive), '--json'], text=True)
rows = [json.loads(line) for line in log.splitlines()]
with tempfile.TemporaryDirectory(prefix='nab-pack-bench-') as tmp:
    root = Path(tmp)
    fragments = []
    for row in rows:
        path = root / f"{row['index']}.record"
        command([new, 'record', archive, row['index'], path])
        fragments.append(path)
    result = {'platform': platform.platform(), 'archive_sha256': sha(archive), 'records': len(rows),
              'runs': a.runs, 'script_sha256': sha(Path(__file__)),
              'method': 'same record fragments; fresh output files; warm cache; CLI wall time includes validation and durability; SHA-256 verification outside timing; binary formats differ in archive framing',
              'binaries': {name: {'path': str(binary), 'sha256': sha(binary)} for name, binary in [('old', old), ('new', new)]}, 'trials': []}
    for trial in range(a.runs):
        for name, binary in ([('old', old), ('new', new)] if trial % 2 == 0 else [('new', new), ('old', old)]):
            packed = root / f'{name}-{trial}.nab'
            elapsed_pack, _ = command([binary, 'pack', packed, *fragments])
            command([binary, 'verify', packed])
            recovered = root / f'{name}-{trial}-recovered.nab'
            elapsed_recover, _ = command([binary, 'recover', packed, recovered])
            command([binary, 'verify', recovered])
            if packed.read_bytes() != recovered.read_bytes():
                raise RuntimeError('recovered records differ from packed records')
            log2 = subprocess.check_output([str(binary), 'log', str(packed), '--json'], text=True)
            if [json.loads(line)['sha256'] for line in log2.splitlines()] != [r['sha256'] for r in rows]:
                raise RuntimeError('content hashes differ')
            result['trials'].append({'trial': trial, 'binary': name, 'pack_ms': elapsed_pack,
                                     'recover_ms': elapsed_recover, 'archive_bytes': packed.stat().st_size})
    result['medians'] = {name: {metric: statistics.median(t[metric] for t in result['trials'] if t['binary'] == name)
                                 for metric in ('pack_ms', 'recover_ms')} for name in ('old', 'new')}
    a.output.write_text(json.dumps(result, indent=2) + '\n')
    print(json.dumps(result['medians'], indent=2))