#!/usr/bin/env python3
#
# Runs a performance test on some tasks to measure parse-time performance.

import datetime
import json
import os
import subprocess
import time

from third_party.python import colorlog
from third_party.python.absl import app, flags

handler = colorlog.StreamHandler()
handler.setFormatter(colorlog.ColoredFormatter('%(log_color)s%(levelname)s: %(message)s'))
log = colorlog.getLogger(__name__)
log.addHandler(handler)
log.propagate = False  # Needed to stop double logging?

flags.DEFINE_string('plz', 'plz', 'Binary to run to invoke plz')
flags.DEFINE_integer('num_threads', 10, 'Number of parallel threads to give plz')
flags.DEFINE_string('output', 'results.json', 'File to write results to')
flags.DEFINE_string('revision', 'unknown', 'Git revision')
flags.DEFINE_integer('number', 5, 'Number of times to run test')
flags.DEFINE_string('root', 'tree', 'Directory to run in')
FLAGS = flags.FLAGS


def plz() -> list:
    """Returns the plz invocation for a subprocess."""
    return [
        "/usr/bin/time",
        "-f", "%e %M",
        FLAGS.plz,
        '--repo_root', FLAGS.root,
        '--num_threads', str(FLAGS.num_threads),
        'query', 'alltargets',
    ]


def run(i: int):
    """Run once and return the length of time taken."""
    log.info('Run %d of %d', i + 1, FLAGS.number)
    duration, mem = parse_time_output(subprocess.run(plz(), check=True, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE).stderr.decode("utf-8"))
    log.info('Complete in %0.2fs, using %d KB', duration, mem)
    return duration, mem

def parse_time_output(output):
    parts = output.split(" ")
    return float(parts[0].strip()), int(parts[1].strip())

def main(argv):
    FLAGS.root = os.path.abspath(FLAGS.root)
    results = [run(i) for i in range(FLAGS.number)]

    time_results = [time for time, _ in results ]
    mem_results = [mem for _, mem in results ]

    time_results.sort()
    median_time = time_results[len(time_results)//2]

    mem_results.sort()
    median_mem = mem_results[len(mem_results)//2]

    log.info('Complete, median time: %0.2fs, median mem: %0.2f KB', median_time, median_mem)
    log.info('Running again to generate profiling info')
    profile_file = os.path.join(os.getcwd(), 'plz.prof')
    subprocess.check_call(plz() + ['--profile_file', profile_file], stdout=subprocess.DEVNULL)
    log.info('Generating results')
    with open(FLAGS.output, 'w') as f:
        json.dump({
            'revision': FLAGS.revision,
            'timestamp': datetime.datetime.now().isoformat(),
            'parse': {
                'raw': time_results,
                'median': median_time,
            },
            'mem': {
                'raw': mem_results,
                'median': median_mem,
            }
        }, f)
        f.write('\n')


if __name__ == '__main__':
    app.run(main)

@@ -1555,16 +1555,17 @@
 n, mem%0A%0A
+%0A
 def pars
@@ -1676,16 +1676,369 @@
 rip())%0A%0A
+%0Adef read_cpu_info():%0A    %22%22%22Return the CPU model number & number of CPUs.%22%22%22%0A    try:%0A        with open('/proc/cpuinfo') as f:%0A            models = %5Bline%5Bline.index(':')+2:%5D for line in f if line.startswith('model name')%5D%0A        return models%5B0%5D.strip(), len(models)%0A    except:%0A        log.exception('Failed to read CPU info')%0A        return '', 0%0A%0A%0A
 def main
@@ -2695,16 +2695,91 @@
 EVNULL)%0A
+    log.info('Reading CPU info')%0A    cpu_model, num_cpus = read_cpu_info()%0A
     log.
@@ -3175,16 +3175,16 @@
 an_mem,%0A
-
         
@@ -3188,16 +3188,122 @@
        %7D
+,%0A            'cpu': %7B%0A                'model': cpu_model,%0A                'num': num_cpus,%0A            %7D,
 %0A       

Add cpu-counting code to benchmark (#1984)