scripts/benchmarks/benchmark_reader_open.py

"""Opt-in, read-only browser benchmark for pinned samples and random Reader opens."""
import argparse
import json
import statistics
import time
from datetime import datetime, timezone
from pathlib import Path
from urllib.parse import urlencode

from playwright.sync_api import sync_playwright


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--base-url', required=True)
    parser.add_argument('--path-prefix', default='')
    parser.add_argument('--samples', type=Path, required=True,
                        help='JSON array of {name, params}: pinned Reader query parameters')
    parser.add_argument('--repeat', type=int, default=2)
    parser.add_argument('--random-attempts', type=int, default=0)
    parser.add_argument('--timeout', type=float, default=45)
    parser.add_argument('--browser', choices=['chromium', 'firefox'], default='chromium')
    parser.add_argument('--output', type=Path, required=True)
    parser.add_argument('--max-ready-seconds', type=float,
                        help='Optional explicit acceptance budget; failures always exit nonzero')
    args = parser.parse_args()
    if args.repeat < 1 or args.timeout <= 0 or args.random_attempts < 0:
        parser.error('repeat/timeout must be positive and random-attempts nonnegative')
    samples = json.loads(args.samples.read_text())
    if not isinstance(samples, list) or any(not isinstance(item, dict) or not isinstance(item.get('params'), dict) or not item.get('name') for item in samples):
        parser.error('samples must be an array of named Reader parameter objects')
    base = args.base_url.rstrip('/')
    prefix = args.path_prefix.rstrip('/')
    report = {'started_at': datetime.now(timezone.utc).isoformat(), 'base_url': base,
              'path_prefix': prefix, 'samples': samples, 'timeout_seconds': args.timeout,
              'cache_policy': 'default HTTP cache, active Service Worker, no routed requests',
              'results': []}
    rows = report['results']

    def save():
        args.output.write_text(json.dumps(report, ensure_ascii=False, indent=2))

    with sync_playwright() as playwright:
        browser = getattr(playwright, args.browser).launch(headless=True, args=['--no-sandbox'] if args.browser == 'chromium' else [])
        report['browser_engine'] = args.browser
        report['browser'] = browser.version
        context = browser.new_context(viewport={'width': 390, 'height': 844})
        context.add_init_script('''window.__readerStages=[];
          new MutationObserver(()=>{
            const item={at:performance.now(),phase:document.documentElement.dataset.readerPhase,
              status:document.querySelector('#status')?.textContent};
            const previous=__readerStages.at(-1);
            if(!previous||previous.phase!==item.phase||previous.status!==item.status) __readerStages.push(item);
          }).observe(document,{subtree:true,childList:true,attributes:true,characterData:true});''')
        page = context.new_page()
        page_errors = []
        page.on('pageerror', lambda error: page_errors.append(str(error)))
        page.set_default_timeout(args.timeout * 1000)
        requests = {}
        cdp = context.new_cdp_session(page) if args.browser == 'chromium' else None
        if cdp:
            cdp.send('Network.enable')

        def sent(event):
            headers = event['request'].get('headers', {})
            requests[event['requestId']] = {'url': event['request']['url'], 'started': event['timestamp'],
                                          'method': event['request']['method'],
                                          'range': next((value for name, value in headers.items() if name.lower() == 'range'), None)}

        def received(event):
            row = requests.get(event['requestId'])
            if row is not None:
                response = event['response']
                row.update(status=response['status'], headers_ms=round((event['timestamp']-row['started'])*1000, 2),
                           server_timing={k:v for k,v in response.get('headers', {}).items() if k.lower()=='server-timing'},
                           transport_timing=response.get('timing'), from_sw=response.get('fromServiceWorker', False),
                           protocol=response.get('protocol'), connection_id=response.get('connectionId'),
                           connection_reused=response.get('connectionReused'), remote_address=response.get('remoteIPAddress'),
                           from_disk_cache=response.get('fromDiskCache', False))

        def data_received(event):
            row = requests.get(event['requestId'])
            if row is not None:
                at = round((event['timestamp'] - row['started']) * 1000, 2)
                row.setdefault('first_data_ms', at)
                if 'last_data_ms' in row:
                    row['max_data_gap_ms'] = max(row.get('max_data_gap_ms', 0), round(at - row['last_data_ms'], 2))
                row['last_data_ms'] = at
                row['received_bytes'] = row.get('received_bytes', 0) + event.get('dataLength', 0)

        def finished(event):
            row = requests.get(event['requestId'])
            if row is not None:
                row.update(total_ms=round((event['timestamp']-row['started'])*1000, 2),
                           bytes=event.get('encodedDataLength'), error=event.get('errorText'))

        if cdp:
            cdp.on('Network.requestWillBeSent', sent)
            cdp.on('Network.responseReceived', received)
            cdp.on('Network.dataReceived', data_received)
            cdp.on('Network.loadingFinished', finished)
            cdp.on('Network.loadingFailed', finished)

        def home():
            page.goto(base + prefix + '/', wait_until='domcontentloaded')
            page.wait_for_function('()=>typeof navigateToReader==="function" && !!navigator.serviceWorker.controller')
            page.wait_for_timeout(2000)

        try:
            home()
            selections = [('random', None, attempt) for attempt in range(args.random_attempts)]
            selections += [(item['name'], item['params'], attempt) for item in samples for attempt in range(args.repeat)]
            for name, params, attempt in selections:
                requests.clear()
                started = time.monotonic()
                row = {'sample': name, 'attempt': attempt}
                return_url = page.url
                reset_host = False
                try:
                    if params is None:
                        page.evaluate('()=>{randomTxt();}')
                    else:
                        page.evaluate('url=>navigateToReader(url)', base + prefix + '/static/reader.html?' + urlencode(params))
                    overlay = page.locator('iframe.reader-overlay')
                    overlay.wait_for()
                    row['iframe_seconds'] = round(time.monotonic()-started, 3)
                    frame = overlay.element_handle().content_frame()
                    frame.locator('.reader-toolbar').wait_for()
                    row['shell_seconds'] = round(time.monotonic()-started, 3)
                    frame.wait_for_function('()=>["ready","failed"].includes(document.documentElement.dataset.readerPhase)')
                    row['ready_seconds'] = round(time.monotonic()-started, 3)
                    row.update(frame.evaluate(r'''()=>({phase:document.documentElement.dataset.readerPhase,
                      error_code:document.querySelector('#content')?.dataset.errorCode,
                      reader_mode:document.querySelector('#content')?.dataset.mode,
                      stages:window.__readerStages,
                      resources:performance.getEntriesByType('resource').map(e=>e.toJSON()),
                      assets:performance.getEntriesByType('resource').filter(e=>/\.(js|mjs|css)(\?|$)/.test(e.name)).map(e=>e.name)})'''))
                    if row['phase'] != 'ready':
                        raise RuntimeError('Reader failed: ' + str(row.get('error_code')))
                    if row.get('reader_mode') in ('pdf', 'pdf-pages', 'image-pages'):
                        target_page = frame.locator('#page-number').input_value()
                        frame.locator(f'.reader-page[data-page="{target_page}"] img.ready, '
                                      f'.reader-page[data-page="{target_page}"] canvas.ready').wait_for()
                        row['first_surface_seconds'] = round(time.monotonic()-started, 3)
                    if row['reader_mode'] in ('audio', 'video'):
                        frame.wait_for_function('()=>document.querySelector("audio,video")?.readyState>=1')
                        row['media_metadata_seconds'] = round(time.monotonic()-started, 3)
                    row['success'] = True
                    frame.locator('#back').click()
                    overlay.wait_for(state='detached', timeout=5000)
                    page.wait_for_function('url=>location.pathname===new URL(url).pathname', arg=return_url, timeout=5000)
                except Exception as error:
                    row.update(success=False, error=str(error), elapsed_seconds=round(time.monotonic()-started, 3))
                    reset_host = True
                row['network'] = list(requests.values())
                rows.append(row)
                print(json.dumps({k:v for k,v in row.items() if k not in ('network','stages','assets','resources')}, ensure_ascii=False), flush=True)
                save()
                if reset_host:
                    # Closing only the iframe leaves an overlay history entry that
                    # a later Back can reopen. Reset this benchmark's failed trial.
                    page.evaluate('''url=>{
                      closeReaderOverlay();
                      sessionStorage.removeItem('reader-navigation-current');
                      history.replaceState(null,'',url);
                    }''', base + prefix + '/')
                    home()
                page.wait_for_timeout(500)
        except Exception as error:
            row = {'sample': 'search-host', 'success': False, 'error': str(error),
                   'page_errors': page_errors, 'network': list(requests.values())}
            rows.append(row)
            print(json.dumps(row), flush=True)
        finally:
            context.close()
            browser.close()
            save()
    report['summary'] = {}
    for name in {row['sample'] for row in rows}:
        selected = [row for row in rows if row['sample'] == name]
        times = [max(row['ready_seconds'], row.get('first_surface_seconds', 0), row.get('media_metadata_seconds', 0))
                 for row in selected if row.get('success')]
        report['summary'][name] = {'runs': len(selected), 'failures': sum(not row.get('success') for row in selected),
                                  'median_usable_seconds': statistics.median(times) if times else None,
                                  'max_usable_seconds': max(times) if times else None}
    save()
    print(json.dumps(report['summary'], ensure_ascii=False), flush=True)
    return int(any(not row.get('success') or args.max_ready_seconds is not None and
                   max(row.get('ready_seconds', float('inf')), row.get('media_metadata_seconds', 0)) > args.max_ready_seconds for row in rows))


if __name__ == '__main__':
    raise SystemExit(main())