#!/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 objgraph
import tracemalloc
from datetime import datetime, timezone, timedelta
from fractions import Fraction
from decimal import Decimal
from collections import namedtuple, OrderedDict

def import_cbor2():
    # Similar hack to that used in tests/conftest to get separate C and Python
    # implementations
    import cbor2
    import cbor2.types
    import cbor2.encoder
    import cbor2.decoder
    class Module(object):
        # Mock module class
        pass
    py_cbor2 = Module()
    for source in (cbor2.types, cbor2.encoder, cbor2.decoder):
        for name in dir(source):
            setattr(py_cbor2, name, getattr(source, name))
    return cbor2, py_cbor2

c_cbor2, py_cbor2 = 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),
]

Leaks = namedtuple('Leaks', ('count', 'comparison'))
Tests = namedtuple('Test', ('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=102, 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:
        return '%d objs (/%d)' % (
            sum(leak[-1] for leak in result.objgraph.comparison),
            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 '%d bytes (/%d)' % (
            result.malloc.comparison[0].size_diff, 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 = py_cbor2.dumps(value, **kwargs)
        results[name] = Result(
            encoding=test(lambda: c_cbor2.dumps(value, **kwargs)),
            decoding=test(lambda: c_cbor2.loads(encoded)),
            roundtrip=test(lambda: c_cbor2.loads(c_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()
