#!/usr/bin/env python

"""
This is a crude script for detecting reference leaks in the C-based cbor2
implementation. It is by no means fool-proof and won't pick up all possible ref
leaks, but it is a reasonable "confidence test" that things aren't horribly
wrong. The script assumes you're in an environment with objgraph and cbor2
installed.

The script outputs a nicely formatted table of the tests run, and the number of
"extra" objects that existed after the tests (indicating a ref-leak), or "-" if
no extra objects existed. The ideal output is obviously "-" in all rows.
"""

import sys
import tracemalloc
from collections import OrderedDict, namedtuple
from datetime import datetime, timedelta, timezone
from decimal import Decimal
from fractions import Fraction

import objgraph

import cbor2

UTC = timezone.utc

TEST_VALUES = [
    # label, kwargs, value
    ("None", {}, None),
    ("10e0", {}, 1),
    ("10e12", {}, 1000000000000),
    ("10e29", {}, 100000000000000000000000000000),
    ("-10e0", {}, -1),
    ("-10e12", {}, -1000000000000),
    ("-10e29", {}, -100000000000000000000000000000),
    ("float1", {}, 1.0),
    ("float2", {}, 3.8),
    ("str", {}, "foo"),
    ("bigstr", {}, "foobarbaz " * 1000),
    ("bytes", {}, b"foo"),
    ("bigbytes", {}, b"foobarbaz\x00" * 1000),
    ("datetime", {"timezone": UTC}, datetime(2019, 5, 9, 22, 4, 5, 123456)),
    ("decimal", {}, Decimal("1.1")),
    ("fraction", {}, Fraction(1, 5)),
    ("intlist", {}, [1, 2, 3]),
    ("bigintlist", {}, [1, 2, 3] * 1000),
    ("strlist", {}, ["foo", "bar", "baz"]),
    ("bigstrlist", {}, ["foo", "bar", "baz"] * 1000),
    ("dict", {}, {"a": 1, "b": 2, "c": 3}),
    ("bigdict", {}, {"a" * i: i for i in range(1000)}),
    ("set", {}, {1, 2, 3}),
    ("bigset", {}, set(range(1000))),
    ("bigdictlist", {}, [{"a" * i: i for i in range(100)}] * 100),
    (
        "objectdict",
        {"timezone": UTC},
        {"name": "Foo", "species": "cat", "dob": datetime(2013, 5, 20), "weight": 4.1},
    ),
    (
        "objectdictlist",
        {"timezone": UTC},
        [{"name": "Foo", "species": "cat", "dob": datetime(2013, 5, 20), "weight": 4.1}] * 100,
    ),
    ("tag", {}, cbor2.CBORTag(1, 1)),
    ("nestedtag", {}, {cbor2.CBORTag(1, 1): 1}),
    ("longstr_128k", {}, "x" * 131072),
    ("longstr_multi_utf8", {}, ("a" * 65535 + "€") * 2),
]

Leaks = namedtuple("Leaks", ("count", "comparison"))
Tests = namedtuple("Tests", ("objgraph", "malloc"))
Result = namedtuple("Result", ("encoding", "decoding", "roundtrip"))


peak = {}


def growth():
    return objgraph.growth(limit=None, peak_stats=peak)


def test_malloc(op):
    count = 0
    start = datetime.now()
    # NOTE: Filter pointing to the op() line in the loop below, because we're
    # only interested in memory allocated by that line. Naturally, if this file
    # is edited, the lineno parameter below must be adjusted!
    only_op = tracemalloc.Filter(True, __file__, lineno=119, all_frames=True)
    tracemalloc.start(10)
    try:
        # Perform a pre-run of op so that any one-time memory allocation
        # (module imports, etc.) don't affect the later diffs
        op()
        before = tracemalloc.take_snapshot().filter_traces([only_op])
        while True:
            count += 1
            op()
            if datetime.now() - start > timedelta(seconds=0.2):
                break
        after = tracemalloc.take_snapshot().filter_traces([only_op])
        diff = after.compare_to(before, "traceback")
        diff = [entry for entry in diff if entry.size_diff > 0]
        return count, diff
    finally:
        tracemalloc.stop()


def test_objgraph(op):
    count = 0
    start = datetime.now()
    # See notes above
    op()
    growth()
    while True:
        count += 1
        op()
        if datetime.now() - start > timedelta(seconds=0.2):
            break
    return count, growth()


def test(op):
    return Tests(Leaks(*test_objgraph(op)), Leaks(*test_malloc(op)))


def format_leaks(result):
    if result.objgraph.comparison:
        num_objs = sum(leak[-1] for leak in result.objgraph.comparison)
        return f"{num_objs} objs (/{result.objgraph.count})"
    elif result.malloc.comparison and (
        result.malloc.count < result.malloc.comparison[0].size_diff
    ):
        # Running the loop always results in *some* memory allocation, but as
        # long as the bytes allocated are less than the number of loops it's
        # unlikely to be an actual leak
        return f"{result.malloc.comparison[0].size_diff} bytes (/{result.malloc.count})"
    else:
        return "-"


def output_table(results):
    # Build table content
    head = ("Test", "Encoding", "Decoding", "Round-trip")
    rows = [head] + [
        (
            label,
            format_leaks(result.encoding),
            format_leaks(result.decoding),
            format_leaks(result.roundtrip),
        )
        for label, result in results.items()
    ]

    # Format table output
    cols = zip(*rows)
    col_widths = [max(len(row) for row in col) for col in cols]
    sep = "".join(
        (
            "+-",
            "-+-".join("-" * width for width in col_widths),
            "-+",
        )
    )
    print(sep)
    print(
        "".join(
            (
                "| ",
                " | ".join(
                    "{value:<{width}}".format(value=value, width=width)
                    for value, width in zip(head, col_widths)
                ),
                " |",
            )
        )
    )
    print(sep)
    for row in rows[1:]:
        print(
            "".join(
                (
                    "| ",
                    " | ".join(
                        "{value:<{width}}".format(value=value, width=width)
                        for value, width in zip(row, col_widths)
                    ),
                    " |",
                )
            )
        )
    print(sep)
    print()
    print(
        """\
There *will* be false positives in the table above. Ignore leaks involving a
tiny number of objects (e.g. 1) or a small number of bytes (e.g. < 8Kb) as such
allocations are quite normal.

In the case of a ref-leak of an object that can reference others (lists, sets,
dicts, or anything with a __dict__), expect to see 100s or 1000s of "objs"
leaked. In the case of a ref-leak of a simple object (int, str, bytes, etc.),
expect to see a few hundred Kb allocated.

If leaks occur across the board, it's likely to be in something universal like
dump/load. If it's restricted to a type, check the encoding and decoding
methods for that type.
"""
    )


def main():
    results = OrderedDict()
    sys.stderr.write("Testing")
    sys.stderr.flush()
    for name, kwargs, value in TEST_VALUES:
        encoded = cbor2.dumps(value, **kwargs)
        results[name] = Result(
            encoding=test(lambda: cbor2.dumps(value, **kwargs)),
            decoding=test(lambda: cbor2.loads(encoded)),
            roundtrip=test(lambda: cbor2.loads(cbor2.dumps(value, **kwargs))),
        )
        sys.stderr.write(".")
        sys.stderr.flush()
    sys.stderr.write("\n")
    sys.stderr.write("\n")
    output_table(results)
    sys.stderr.write("\n")


if __name__ == "__main__":
    main()
