#!/usr/bin/env python

# usage:
# leaktest.py "default" iteration-count collection-uuid
# leaktest.py "workaround" iteration-count collection-uuid

import arvados
import os
import resource
import subprocess
import sys

_, workaround, iterations, uuid = sys.argv

api = arvados.api('v1')
if workaround == 'workaround':
    print "Using block_manager workaround"
    block_manager = arvados.collection.CollectionReader(
        uuid, api_client=api)._my_block_manager()
else:
    print "Not using block_manager workaround"
    block_manager = None

for i in range(int(iterations)):
    print "{}...".format(i+1),
    sys.stdout.flush()
    cr = arvados.collection.CollectionReader(
        uuid, api_client=api, block_manager=block_manager)
    for fn in cr:
        f = cr.open(fn)
        f.read()
        f.close()
print ""

print "Max RSS: {}".format(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss)
conn = 0
for line in subprocess.check_output(['netstat', '-pane'], stderr=open(os.devnull, 'w')).splitlines():
    if line.find(str(os.getpid())+"/python") >= 0:
        conn += 1
print "Network connections still open: {}".format(conn)
