diff --git a/bitbake/bin/bitbake-hashclient b/bitbake/bin/bitbake-hashclient index 494f17592a..7c57866d36 100755 --- a/bitbake/bin/bitbake-hashclient +++ b/bitbake/bin/bitbake-hashclient @@ -56,25 +56,24 @@ def main(): nonlocal missed_hashes nonlocal max_time - client = hashserv.create_client(args.address) + with hashserv.create_client(args.address) as client: + for i in range(args.requests): + taskhash = hashlib.sha256() + taskhash.update(args.taskhash_seed.encode('utf-8')) + taskhash.update(str(i).encode('utf-8')) - for i in range(args.requests): - taskhash = hashlib.sha256() - taskhash.update(args.taskhash_seed.encode('utf-8')) - taskhash.update(str(i).encode('utf-8')) + start_time = time.perf_counter() + l = client.get_unihash(METHOD, taskhash.hexdigest()) + elapsed = time.perf_counter() - start_time - start_time = time.perf_counter() - l = client.get_unihash(METHOD, taskhash.hexdigest()) - elapsed = time.perf_counter() - start_time + with lock: + if l: + found_hashes += 1 + else: + missed_hashes += 1 - with lock: - if l: - found_hashes += 1 - else: - missed_hashes += 1 - - max_time = max(elapsed, max_time) - pbar.update() + max_time = max(elapsed, max_time) + pbar.update() max_time = 0 found_hashes = 0 @@ -152,9 +151,8 @@ def main(): func = getattr(args, 'func', None) if func: - client = hashserv.create_client(args.address) - - return func(args, client) + with hashserv.create_client(args.address) as client: + return func(args, client) return 0 diff --git a/bitbake/lib/bb/asyncrpc/client.py b/bitbake/lib/bb/asyncrpc/client.py index fa042bbe87..dcbe7e5762 100644 --- a/bitbake/lib/bb/asyncrpc/client.py +++ b/bitbake/lib/bb/asyncrpc/client.py @@ -126,6 +126,12 @@ class AsyncClient(object): {'ping': {}} ) + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc_value, traceback): + await self.close() + class Client(object): def __init__(self): @@ -176,3 +182,10 @@ class Client(object): if sys.version_info >= (3, 6): self.loop.run_until_complete(self.loop.shutdown_asyncgens()) self.loop.close() + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + self.close() + return False diff --git a/bitbake/lib/prserv/serv.py b/bitbake/lib/prserv/serv.py index c686b2065c..0db6ebc707 100644 --- a/bitbake/lib/prserv/serv.py +++ b/bitbake/lib/prserv/serv.py @@ -344,9 +344,9 @@ def auto_shutdown(): def ping(host, port): from . import client - conn = client.PRClient() - conn.connect_tcp(host, port) - return conn.ping() + with client.PRClient() as conn: + conn.connect_tcp(host, port) + return conn.ping() def connect(host, port): from . import client