#!/usr/bin/env python3
"""Warm (or hammer) a cgit instance running behind Varnish with ESI.

Fragments are shared by every repository that holds the same objects, so the
diff rendered for one fork's commit is the diff every other fork will be
served. Crawling one repository therefore preloads the expensive half of the
page for all of them; the per-fork frames stay cheap and are rendered on
demand.

Two modes:

  crawl    Breadth-first from a starting URL, following links the way a
           scraper would. Commit and diff pages are queued ahead of
           everything else, since those are the pages that carry fragments.

  commits  Walk a single repository's log pages and fetch every commit they
           name. Denser than crawling, but cgit paginates the log by walking
           the revision list from the start every time, so this gets slower
           the deeper it goes. Fine for the first few thousand commits.

  oids     Read object ids (one 40-hex id per line) and fetch a commit page
           for each. This is the one for a bulk preload: it has no
           pagination to walk, so it stays at full speed all the way
           through a large history.

The reported hit rate counts only the pages the crawler asks for. ESI
fragments are fetched by Varnish itself, and a client cannot see whether
they hit; in the steady state each page is one frame miss plus one fragment
hit, so a healthy run shows a low client-side hit rate and roughly half the
backend requests you would otherwise expect. Read varnishstat for the truth.

Examples:

    # bulk preload every commit, fed straight from git
    git -C /srv/git/linux.git rev-list --all |
        cgit-crawl.py https://kernel.varnish.org/ -m oids \\
            -r pub/scm/linux/kernel/git/torvalds/linux.git -c 32

    # preload the first few thousand commits without needing the repo
    cgit-crawl.py https://kernel.varnish.org/ -m commits \\
        -r pub/scm/linux/kernel/git/torvalds/linux.git -c 32 -n 5000

    # behave like a scraper and see what the cache does
    cgit-crawl.py https://kernel.varnish.org/ -m crawl -c 16 -n 5000
"""

import argparse
import http.client
import queue
import re
import ssl
import sys
import threading
import time
import urllib.parse
from html.parser import HTMLParser

# Pages that cost the backend a lot and teach us nothing about fragments.
SKIP = re.compile(
    r"/(snapshot|patch|rawdiff|atom|blame|plain|stats)/"
    r"|[?&]p=(patch|atom)"
    r"|\.(tar|gz|bz2|xz|zip)$"
)
# The pages that actually contain an ESI fragment.
FRAGMENT_BEARING = re.compile(r"/(commit|diff)/")


class LinkParser(HTMLParser):
    def __init__(self):
        super().__init__(convert_charrefs=True)
        self.hrefs = []

    def handle_starttag(self, tag, attrs):
        if tag == "a":
            for k, v in attrs:
                if k == "href" and v:
                    self.hrefs.append(v)


class Stats:
    def __init__(self):
        self.lock = threading.Lock()
        self.req = 0
        self.hit = 0
        self.miss = 0
        self.bytes = 0
        self.errors = 0
        self.status = {}
        self.start = time.time()

    def record(self, status, cache, nbytes):
        with self.lock:
            self.req += 1
            self.bytes += nbytes
            self.status[status] = self.status.get(status, 0) + 1
            if cache == "hit":
                self.hit += 1
            elif cache == "miss":
                self.miss += 1

    def error(self):
        with self.lock:
            self.errors += 1

    def snapshot(self):
        with self.lock:
            el = max(time.time() - self.start, 1e-6)
            return dict(
                req=self.req, hit=self.hit, miss=self.miss, bytes=self.bytes,
                errors=self.errors, elapsed=el, rate=self.req / el,
                status=dict(self.status),
            )


class Fetcher:
    """One persistent connection, reused across requests by a single thread."""

    def __init__(self, base, insecure):
        self.scheme, self.host = base
        self.insecure = insecure
        self.conn = None

    def _connect(self):
        if self.scheme == "https":
            ctx = ssl.create_default_context()
            if self.insecure:
                ctx.check_hostname = False
                ctx.verify_mode = ssl.CERT_NONE
            self.conn = http.client.HTTPSConnection(self.host, timeout=60,
                                                    context=ctx)
        else:
            self.conn = http.client.HTTPConnection(self.host, timeout=60)

    def get(self, path):
        for attempt in (1, 2):
            try:
                if self.conn is None:
                    self._connect()
                self.conn.request("GET", path, headers={
                    "Host": self.host,
                    "User-Agent": "cgit-crawl/1.0",
                    "Accept-Encoding": "gzip",
                    "Connection": "keep-alive",
                })
                resp = self.conn.getresponse()
                body = resp.read()
                return resp.status, resp.getheader("X-Cache", ""), body
            except Exception:
                # Stale keep-alive, or the server closed on us. Reconnect once.
                try:
                    if self.conn:
                        self.conn.close()
                except Exception:
                    pass
                self.conn = None
                if attempt == 2:
                    raise
        raise AssertionError("unreachable")

    def close(self):
        if self.conn:
            try:
                self.conn.close()
            except Exception:
                pass


def normalize(base_path, href, host):
    """Resolve href against base_path, returning a path or None if off-site."""
    if href.startswith("#") or href.startswith("mailto:"):
        return None
    u = urllib.parse.urljoin(base_path, href)
    p = urllib.parse.urlsplit(u)
    if p.scheme and p.scheme not in ("http", "https"):
        return None
    if p.netloc and p.netloc != host:
        return None
    path = p.path or "/"
    if p.query:
        path += "?" + p.query
    return path


def crawl(args, base, stats, stop):
    """Breadth-first, with commit and diff pages jumping the queue."""
    host = base[1]
    hi, lo = queue.Queue(), queue.Queue()
    seen, seen_lock = set(), threading.Lock()

    start = urllib.parse.urlsplit(args.url).path or "/"
    seen.add(start)
    lo.put(start)

    def take():
        for q in (hi, lo):
            try:
                return q.get_nowait()
            except queue.Empty:
                continue
        return None

    def worker():
        f = Fetcher(base, args.insecure)
        try:
            while not stop.is_set():
                if args.limit and stats.snapshot()["req"] >= args.limit:
                    return
                path = take()
                if path is None:
                    time.sleep(0.05)
                    if hi.empty() and lo.empty():
                        return
                    continue
                try:
                    status, cache, body = f.get(path)
                except Exception:
                    stats.error()
                    continue
                stats.record(status, cache, len(body))
                if status != 200 or b"<a" not in body[:200000]:
                    continue
                try:
                    text = body.decode("utf-8", "replace")
                except Exception:
                    continue
                p = LinkParser()
                try:
                    p.feed(text)
                except Exception:
                    pass
                for href in p.hrefs:
                    nxt = normalize(path, href, host)
                    if not nxt or SKIP.search(nxt):
                        continue
                    if not args.tree and "/tree/" in nxt:
                        continue
                    with seen_lock:
                        if nxt in seen:
                            continue
                        seen.add(nxt)
                    (hi if FRAGMENT_BEARING.search(nxt) else lo).put(nxt)
        finally:
            f.close()

    run_workers(args, worker, stats, stop)


def commits(args, base, stats, stop):
    """Walk log pages of one repository and fetch every commit they name."""
    repo = args.repo.strip("/")
    oid_re = re.compile(r"/commit/\?(?:[^']*&amp;)?id=([0-9a-f]{40})")
    work = queue.Queue(maxsize=args.concurrency * 64)
    done = threading.Event()

    def producer():
        f = Fetcher(base, args.insecure)
        ofs, empty_streak = 0, 0
        try:
            while not stop.is_set():
                if args.limit and stats.snapshot()["req"] >= args.limit:
                    break
                path = "/%s/log/?ofs=%d" % (repo, ofs)
                try:
                    status, cache, body = f.get(path)
                except Exception:
                    stats.error()
                    break
                stats.record(status, cache, len(body))
                if status != 200:
                    break
                oids = oid_re.findall(body.decode("utf-8", "replace"))
                uniq = list(dict.fromkeys(oids))
                if not uniq:
                    empty_streak += 1
                    if empty_streak >= 2:
                        break
                else:
                    empty_streak = 0
                for oid in uniq:
                    if stop.is_set():
                        break
                    work.put("/%s/commit/?id=%s" % (repo, oid))
                ofs += args.log_step
        finally:
            f.close()
            done.set()

    def worker():
        f = Fetcher(base, args.insecure)
        try:
            while not stop.is_set():
                if args.limit and stats.snapshot()["req"] >= args.limit:
                    return
                try:
                    path = work.get(timeout=0.5)
                except queue.Empty:
                    if done.is_set() and work.empty():
                        return
                    continue
                try:
                    status, cache, body = f.get(path)
                except Exception:
                    stats.error()
                    continue
                stats.record(status, cache, len(body))
        finally:
            f.close()

    pt = threading.Thread(target=producer, daemon=True)
    pt.start()
    run_workers(args, worker, stats, stop)
    pt.join(timeout=2)


def oids(args, base, stats, stop):
    """Fetch one commit page per object id read from a file or stdin."""
    repo = args.repo.strip("/")
    valid = re.compile(r"^[0-9a-f]{7,40}$")
    work = queue.Queue(maxsize=args.concurrency * 64)
    done = threading.Event()

    src = sys.stdin if args.oids_from in (None, "-") else open(args.oids_from)

    def producer():
        try:
            for line in src:
                if stop.is_set():
                    break
                oid = line.strip().split()[0] if line.strip() else ""
                if not valid.match(oid):
                    continue
                work.put("/%s/commit/?id=%s" % (repo, oid))
        except Exception:
            pass
        finally:
            done.set()
            if src is not sys.stdin:
                src.close()

    def worker():
        f = Fetcher(base, args.insecure)
        try:
            while not stop.is_set():
                if args.limit and stats.snapshot()["req"] >= args.limit:
                    return
                try:
                    path = work.get(timeout=0.5)
                except queue.Empty:
                    if done.is_set() and work.empty():
                        return
                    continue
                try:
                    status, cache, body = f.get(path)
                except Exception:
                    stats.error()
                    continue
                stats.record(status, cache, len(body))
        finally:
            f.close()

    pt = threading.Thread(target=producer, daemon=True)
    pt.start()
    run_workers(args, worker, stats, stop)
    pt.join(timeout=2)


def run_workers(args, worker, stats, stop):
    threads = [threading.Thread(target=worker, daemon=True)
               for _ in range(args.concurrency)]
    for t in threads:
        t.start()
    last = 0.0
    try:
        while any(t.is_alive() for t in threads):
            time.sleep(0.25)
            now = time.time()
            if args.progress and now - last >= args.progress:
                last = now
                report(stats.snapshot(), end="\r")
    except KeyboardInterrupt:
        stop.set()
    for t in threads:
        t.join(timeout=2)


def human(n):
    for unit in ("B", "KB", "MB", "GB", "TB"):
        if n < 1024 or unit == "TB":
            return "%.1f %s" % (n, unit)
        n /= 1024.0


def report(s, end="\n"):
    served = s["hit"] + s["miss"]
    ratio = (100.0 * s["hit"] / served) if served else 0.0
    sys.stderr.write(
        "%8d req  %6.1f req/s  %5.1f%% frame-hit  %10s  %d err   %s" % (
            s["req"], s["rate"], ratio, human(s["bytes"]), s["errors"],
            end == "\r" and " " or ""))
    sys.stderr.write(end)
    sys.stderr.flush()


def main():
    ap = argparse.ArgumentParser(
        description="Warm a cgit+Varnish ESI cache.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__)
    ap.add_argument("url", help="starting URL, e.g. https://host/")
    ap.add_argument("-m", "--mode", choices=("crawl", "commits", "oids"),
                    default="crawl", help="default: crawl")
    ap.add_argument("-r", "--repo",
                    help="repository path for --mode=commits/oids, e.g. "
                         "pub/scm/linux/kernel/git/torvalds/linux.git")
    ap.add_argument("--oids-from", metavar="FILE",
                    help="read object ids from FILE for --mode=oids "
                         "(default: stdin)")
    ap.add_argument("-c", "--concurrency", type=int, default=16)
    ap.add_argument("-n", "--limit", type=int, default=0,
                    help="stop after roughly this many requests (0 = no limit)")
    ap.add_argument("--log-step", type=int, default=50,
                    help="commits per log page, must match cgit's "
                         "max-commit-count (default: 50)")
    ap.add_argument("--tree", action="store_true",
                    help="also crawl tree pages (off by default: they carry "
                         "no fragments and there are a great many of them)")
    ap.add_argument("--insecure", action="store_true",
                    help="do not verify the TLS certificate")
    ap.add_argument("--progress", type=float, default=1.0,
                    help="seconds between progress lines (0 to disable)")
    args = ap.parse_args()

    if args.mode in ("commits", "oids") and not args.repo:
        ap.error("--mode=%s needs --repo" % args.mode)

    p = urllib.parse.urlsplit(args.url)
    if not p.netloc:
        ap.error("url must be absolute, e.g. https://host/")
    base = (p.scheme or "https", p.netloc)

    stats = Stats()
    stop = threading.Event()
    try:
        if args.mode == "crawl":
            crawl(args, base, stats, stop)
        elif args.mode == "commits":
            commits(args, base, stats, stop)
        else:
            oids(args, base, stats, stop)
    except KeyboardInterrupt:
        stop.set()

    s = stats.snapshot()
    sys.stderr.write("\n")
    report(s)
    if s["status"]:
        sys.stderr.write("status: %s\n" % ", ".join(
            "%s=%d" % kv for kv in sorted(s["status"].items())))
    sys.stderr.write("elapsed: %.1fs\n" % s["elapsed"])


if __name__ == "__main__":
    main()