upstream-layers/bitbake/lib/hashserv/client.py - openbmc - Git at Google

 # Copyright (C) 2019 Garmin Ltd.
 #
 # SPDX-License-Identifier: GPL-2.0-only
 #

 import logging
 import socket
 import asyncio
 import bb.asyncrpc
 import json
 from abc import abstractmethod
 from collections.abc import AsyncIterable
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
 from . import create_async_client

 logger = logging.getLogger("hashserv.client")


 class AsyncQueue(AsyncIterable):
     class Shutdown(Exception):
         pass

     SHUTDOWN_SENTINEL = object()

     def __init__(self, *args, **kwargs):
         self.__queue = asyncio.Queue()
         self.__shutdown = False
         self.__is_done = False

     async def done(self):
         if self.__is_done:
             return
         self.__is_done = True
         await self.__queue.put(self.SHUTDOWN_SENTINEL)

     async def put(self, item):
         if self.__is_done:
             raise self.Shutdown
         await self.__queue.put(item)

     async def get(self):
         if self.__shutdown:
             raise self.Shutdown

         item = await self.__queue.get()
         if item is self.SHUTDOWN_SENTINEL:
             self.__shutdown = True
             raise self.Shutdown

         return item

     def __aiter__(self):
         return self

     async def __anext__(self):
         try:
             return await self.get()
         except self.Shutdown:
             raise StopAsyncIteration


 @dataclass(eq=False, frozen=True)
 class AsyncPipe:
     send_queue: AsyncQueue
     recv_queue: AsyncQueue


 class Stream(AsyncIterable):
     def __init__(self, pipe):
         self._pipe = pipe

     async def done(self):
         await self._pipe.send_queue.done()

     def __aiter__(self):
         return self

     async def __anext__(self):
         try:
             return await self.get_result()
         except AsyncQueue.Shutdown:
             raise StopAsyncIteration

     @abstractmethod
     async def _send_batch_input(self, i):
         raise NotImplementedError("Not implemented")

     @abstractmethod
     async def get_result(self):
         raise NotImplementedError("Not implemented")

     async def batch(self, inputs):
         """
         Does a "batch" process of stream messages. This sends the query
         messages as fast as possible, and simultaneously attempts to read the
         messages back. This helps to mitigate the effects of latency to the
         hash equivalence server be allowing multiple queries to be "in-flight"
         at once

         The input may be a generator or an async generator
         """

         async def get_inputs():
             if isinstance(inputs, AsyncIterable):
                 async for i in inputs:
                     yield i
             else:
                 for i in inputs:
                     yield i

         async def send():
             try:
                 async for i in get_inputs():
                     await self._send_batch_input(i)
             finally:
                 await self.done()

         results = []

         async def recv():
             async for item in self:
                 results.append(item)

         await bb.asyncrpc.TaskGroup.run(send(), recv())
         return results


 class GetUnihashStream(Stream):
     def __init__(self, pipe):
         super().__init__(pipe)

     async def _send_batch_input(self, i):
         method, taskhash = i
         await self.send_query(method, taskhash)

     async def send_query(self, method, taskhash):
         await self._pipe.send_queue.put(f"{method} {taskhash}")

     async def get_result(self):
         r = await self._pipe.recv_queue.get()
         return r if r else None


 class UnihashExistsStream(Stream):
     def __init__(self, pipe):
         super().__init__(pipe)

     async def _send_batch_input(self, i):
         await self.send_query(i)

     async def send_query(self, unihash):
         await self._pipe.send_queue.put(unihash)

     async def get_result(self):
         r = await self._pipe.recv_queue.get()
         return r == "true"


 class GcMarkStream(Stream):
     def __init__(self, pipe, mark):
         super().__init__(pipe)
         self.mark = mark

     async def _send_batch_input(self, i):
         def row_to_dict(row):
             pairs = row.split()
             return dict(zip(pairs[::2], pairs[1::2]))

         await self.send_mark(row_to_dict(i))

     async def send_mark(self, where):
         await self._pipe.send_queue.put(json.dumps({"mark": self.mark, "where": where}))

     async def get_result(self):
         r = await self._pipe.recv_queue.get()
         return json.loads(r)


 class Batch(object):
     def __init__(self, send_queue, recv_queue):
         self.send_queue = send_queue
         self.recv_queue = recv_queue
         self.fill_done = False
         self.send_done = False
         self.cond = asyncio.Condition()
         self.pending = []
         self.sent_count = 0
         self.recv_count = 0
         self.item = None

     async def recv(self, socket):
         while True:
             async with self.cond:
                 await self.cond.wait_for(lambda: self.pending or self.send_done)
                 if not self.pending:
                     if self.send_done:
                         return
                     continue

             m = await socket.recv()
             await self.recv_queue.put(m)

             async with self.cond:
                 self.recv_count += 1
                 self.pending.pop(0)

     async def fill(self):
         async for m in self.send_queue:
             async with self.cond:
                 # Wait for item to be consumed
                 await self.cond.wait_for(lambda: self.item is None)
                 self.item = m
                 self.cond.notify_all()

         async with self.cond:
             self.fill_done = True
             self.cond.notify_all()

     async def send(self, socket):
         # In the event of a restart due to a reconnect, all in-flight
         # messages need to be resent first to keep to result count in sync
         async with self.cond:
             for m in self.pending:
                 await socket.send(m)

         while True:
             async with self.cond:
                 await self.cond.wait_for(
                     lambda: self.item is not None or self.fill_done
                 )
                 if self.item is None:
                     if self.fill_done:
                         self.send_done = True
                         self.cond.notify_all()
                         return
                     continue

                 m = self.item

             await socket.send(m)

             async with self.cond:
                 self.item = None
                 self.pending.append(m)
                 self.sent_count += 1
                 self.cond.notify_all()

     async def stream(self, socket):
         await bb.asyncrpc.TaskGroup.run(self.send(socket), self.recv(socket))

     def check(self):
         if self.sent_count != self.recv_count:
             raise ConnectionError(
                 f"Sent {self.sent_count} messages but only received {self.recv_count}"
             )


 class AsyncClient(bb.asyncrpc.AsyncClient):
     MODE_NORMAL = 0
     MODE_GET_STREAM = 1
     MODE_EXIST_STREAM = 2
     MODE_MARK_STREAM = 3

     def __init__(self, username=None, password=None):
         super().__init__("OEHASHEQUIV", "1.1", logger)
         self.mode = self.MODE_NORMAL
         self.username = username
         self.password = password
         self.saved_become_user = None

     async def setup_connection(self):
         await super().setup_connection()
         self.mode = self.MODE_NORMAL
         if self.username:
             # Save off become user temporarily because auth() resets it
             become = self.saved_become_user
             await self.auth(self.username, self.password)

             if become:
                 await self.become_user(become)

     @asynccontextmanager
     async def send_stream(self, mode):
         send_queue = AsyncQueue()
         recv_queue = AsyncQueue()
         b = Batch(send_queue, recv_queue)

         async def proc():
             await self._set_mode(mode)
             await b.stream(self.socket)

         async def process():
             try:
                 await self._send_wrapper(proc)
             finally:
                 await recv_queue.done()

         # Create background process to process messages
         async with bb.asyncrpc.TaskGroup() as group:
             group.create_task(process())
             group.create_task(b.fill())
             try:
                 yield AsyncPipe(send_queue, recv_queue)
                 b.check()
             except AsyncQueue.Shutdown as e:
                 pass
             finally:
                 await send_queue.done()
                 await recv_queue.done()

     async def invoke(self, *args, skip_mode=False, **kwargs):
         # It's OK if connection errors cause a failure here, because the mode
         # is also reset to normal on a new connection
         if not skip_mode:
             await self._set_mode(self.MODE_NORMAL)
         return await super().invoke(*args, **kwargs)

     async def _set_mode(self, new_mode):
         async def stream_to_normal():
             # Check if already in normal mode (e.g. due to a connection reset)
             if self.mode == self.MODE_NORMAL:
                 return "ok"
             await self.socket.send("END")
             return await self.socket.recv()

         async def normal_to_stream(command):
             r = await self.invoke({command: None}, skip_mode=True)
             if r != "ok":
                 self.check_invoke_error(r)
                 raise ConnectionError(
                     f"Unable to transition to stream mode: Bad response from server {r!r}"
                 )
             self.logger.debug("Mode is now %s", command)

         if new_mode == self.mode:
             return

         self.logger.debug("Transitioning mode %s -> %s", self.mode, new_mode)

         # Always transition to normal mode before switching to any other mode
         if self.mode != self.MODE_NORMAL:
             r = await self._send_wrapper(stream_to_normal)
             if r != "ok":
                 self.check_invoke_error(r)
                 raise ConnectionError(
                     f"Unable to transition to normal mode: Bad response from server {r!r}"
                 )
             self.logger.debug("Mode is now normal")

         if new_mode == self.MODE_GET_STREAM:
             await normal_to_stream("get-stream")
         elif new_mode == self.MODE_EXIST_STREAM:
             await normal_to_stream("exists-stream")
         elif new_mode == self.MODE_MARK_STREAM:
             await normal_to_stream("gc-mark-stream")
         elif new_mode != self.MODE_NORMAL:
             raise Exception("Undefined mode transition {self.mode!r} -> {new_mode!r}")

         self.mode = new_mode

     async def get_unihash(self, method, taskhash):
         async with self.get_unihash_stream() as stream:
             await stream.send_query(method, taskhash)
             return await stream.get_result()

     async def get_unihash_batch(self, args):
         async with self.get_unihash_stream() as stream:
             return await stream.batch(args)

     @asynccontextmanager
     async def get_unihash_stream(self):
         async with self.send_stream(self.MODE_GET_STREAM) as pipe:
             yield GetUnihashStream(pipe)

     async def report_unihash(self, taskhash, method, outhash, unihash, extra={}):
         m = extra.copy()
         m["taskhash"] = taskhash
         m["method"] = method
         m["outhash"] = outhash
         m["unihash"] = unihash
         return await self.invoke({"report": m})

     async def report_unihash_equiv(self, taskhash, method, unihash, extra={}):
         m = extra.copy()
         m["taskhash"] = taskhash
         m["method"] = method
         m["unihash"] = unihash
         return await self.invoke({"report-equiv": m})

     async def get_taskhash(self, method, taskhash, all_properties=False):
         return await self.invoke(
             {"get": {"taskhash": taskhash, "method": method, "all": all_properties}}
         )

     async def unihash_exists(self, unihash):
         async with self.unihash_exists_stream() as stream:
             await stream.send_query(unihash)
             return await stream.get_result()

     async def unihash_exists_batch(self, unihashes):
         async with self.unihash_exists_stream() as stream:
             return await stream.batch(unihashes)

     @asynccontextmanager
     async def unihash_exists_stream(self):
         async with self.send_stream(self.MODE_EXIST_STREAM) as pipe:
             yield UnihashExistsStream(pipe)

     async def get_outhash(self, method, outhash, taskhash, with_unihash=True):
         return await self.invoke(
             {
                 "get-outhash": {
                     "outhash": outhash,
                     "taskhash": taskhash,
                     "method": method,
                     "with_unihash": with_unihash,
                 }
             }
         )

     async def get_stats(self):
         return await self.invoke({"get-stats": None})

     async def reset_stats(self):
         return await self.invoke({"reset-stats": None})

     async def backfill_wait(self):
         return (await self.invoke({"backfill-wait": None}))["tasks"]

     async def remove(self, where):
         return await self.invoke({"remove": {"where": where}})

     async def clean_unused(self, max_age):
         return await self.invoke({"clean-unused": {"max_age_seconds": max_age}})

     async def auth(self, username, token):
         result = await self.invoke({"auth": {"username": username, "token": token}})
         self.username = username
         self.password = token
         self.saved_become_user = None
         return result

     async def refresh_token(self, username=None):
         m = {}
         if username:
             m["username"] = username
         result = await self.invoke({"refresh-token": m})
         if (
             self.username
             and not self.saved_become_user
             and result["username"] == self.username
         ):
             self.password = result["token"]
         return result

     async def set_user_perms(self, username, permissions):
         return await self.invoke(
             {"set-user-perms": {"username": username, "permissions": permissions}}
         )

     async def get_user(self, username=None):
         m = {}
         if username:
             m["username"] = username
         return await self.invoke({"get-user": m})

     async def get_all_users(self):
         return (await self.invoke({"get-all-users": {}}))["users"]

     async def new_user(self, username, permissions):
         return await self.invoke(
             {"new-user": {"username": username, "permissions": permissions}}
         )

     async def delete_user(self, username):
         return await self.invoke({"delete-user": {"username": username}})

     async def become_user(self, username):
         result = await self.invoke({"become-user": {"username": username}})
         if username == self.username:
             self.saved_become_user = None
         else:
             self.saved_become_user = username
         return result

     async def get_db_usage(self):
         return (await self.invoke({"get-db-usage": {}}))["usage"]

     async def get_db_query_columns(self):
         return (await self.invoke({"get-db-query-columns": {}}))["columns"]

     async def gc_status(self):
         return await self.invoke({"gc-status": {}})

     async def gc_mark(self, mark, where):
         """
         Starts a new garbage collection operation identified by "mark". If
         garbage collection is already in progress with "mark", the collection
         is continued.

         All unihash entries that match the "where" clause are marked to be
         kept. In addition, any new entries added to the database after this
         command will be automatically marked with "mark"
         """
         return await self.invoke({"gc-mark": {"mark": mark, "where": where}})

     async def gc_mark_batch(self, mark, rows):
         """
         Similar to `gc-mark`, but accepts a list of "where" key-value pair
         conditions. It utilizes stream mode to mark hashes, which helps reduce
         the impact of latency when communicating with the hash equivalence
         server.
         """
         async with self.gc_mark_stream(mark) as stream:
             results = await stream.batch(rows)

         return {"count": sum(int(r["count"]) for r in results)}

     @asynccontextmanager
     async def gc_mark_stream(self, mark):
         async with self.send_stream(self.MODE_MARK_STREAM) as pipe:
             yield GcMarkStream(pipe, mark)

     async def gc_sweep(self, mark):
         """
         Finishes garbage collection for "mark". All unihash entries that have
         not been marked will be deleted.

         It is recommended to clean unused outhash entries after running this to
         cleanup any dangling outhashes
         """
         return await self.invoke({"gc-sweep": {"mark": mark}})


 class Client(bb.asyncrpc.Client):
     def __init__(self, username=None, password=None):
         self.username = username
         self.password = password

         super().__init__()
         self._add_methods(
             "connect_tcp",
             "connect_websocket",
             "get_unihash",
             "get_unihash_batch",
             "report_unihash",
             "report_unihash_equiv",
             "get_taskhash",
             "unihash_exists",
             "unihash_exists_batch",
             "get_outhash",
             "get_stats",
             "reset_stats",
             "backfill_wait",
             "remove",
             "clean_unused",
             "auth",
             "refresh_token",
             "set_user_perms",
             "get_user",
             "get_all_users",
             "new_user",
             "delete_user",
             "become_user",
             "get_db_usage",
             "get_db_query_columns",
             "gc_status",
             "gc_mark",
             "gc_mark_batch",
             "gc_sweep",
         )

     def _get_async_client(self):
         return AsyncClient(self.username, self.password)
	# Copyright (C) 2019 Garmin Ltd.
	#
	# SPDX-License-Identifier: GPL-2.0-only
	#

	import logging
	import socket
	import asyncio
	import bb.asyncrpc
	import json
	from abc import abstractmethod
	from collections.abc import AsyncIterable
	from contextlib import asynccontextmanager
	from dataclasses import dataclass
	from . import create_async_client

	logger = logging.getLogger("hashserv.client")


	class AsyncQueue(AsyncIterable):
	class Shutdown(Exception):
	pass

	SHUTDOWN_SENTINEL = object()

	def __init__(self, args, *kwargs):
	self.__queue = asyncio.Queue()
	self.__shutdown = False
	self.__is_done = False

	async def done(self):
	if self.__is_done:
	return
	self.__is_done = True
	await self.__queue.put(self.SHUTDOWN_SENTINEL)

	async def put(self, item):
	if self.__is_done:
	raise self.Shutdown
	await self.__queue.put(item)

	async def get(self):
	if self.__shutdown:
	raise self.Shutdown

	item = await self.__queue.get()
	if item is self.SHUTDOWN_SENTINEL:
	self.__shutdown = True
	raise self.Shutdown

	return item

	def __aiter__(self):
	return self

	async def __anext__(self):
	try:
	return await self.get()
	except self.Shutdown:
	raise StopAsyncIteration


	@dataclass(eq=False, frozen=True)
	class AsyncPipe:
	send_queue: AsyncQueue
	recv_queue: AsyncQueue


	class Stream(AsyncIterable):
	def __init__(self, pipe):
	self._pipe = pipe

	async def done(self):
	await self._pipe.send_queue.done()

	def __aiter__(self):
	return self

	async def __anext__(self):
	try:
	return await self.get_result()
	except AsyncQueue.Shutdown:
	raise StopAsyncIteration

	@abstractmethod
	async def _send_batch_input(self, i):
	raise NotImplementedError("Not implemented")

	@abstractmethod
	async def get_result(self):
	raise NotImplementedError("Not implemented")

	async def batch(self, inputs):
	"""
	Does a "batch" process of stream messages. This sends the query
	messages as fast as possible, and simultaneously attempts to read the
	messages back. This helps to mitigate the effects of latency to the
	hash equivalence server be allowing multiple queries to be "in-flight"
	at once

	The input may be a generator or an async generator
	"""

	async def get_inputs():
	if isinstance(inputs, AsyncIterable):
	async for i in inputs:
	yield i
	else:
	for i in inputs:
	yield i

	async def send():
	try:
	async for i in get_inputs():
	await self._send_batch_input(i)
	finally:
	await self.done()

	results = []

	async def recv():
	async for item in self:
	results.append(item)

	await bb.asyncrpc.TaskGroup.run(send(), recv())
	return results


	class GetUnihashStream(Stream):
	def __init__(self, pipe):
	super().__init__(pipe)

	async def _send_batch_input(self, i):
	method, taskhash = i
	await self.send_query(method, taskhash)

	async def send_query(self, method, taskhash):
	await self._pipe.send_queue.put(f"{method} {taskhash}")

	async def get_result(self):
	r = await self._pipe.recv_queue.get()
	return r if r else None


	class UnihashExistsStream(Stream):
	def __init__(self, pipe):
	super().__init__(pipe)

	async def _send_batch_input(self, i):
	await self.send_query(i)

	async def send_query(self, unihash):
	await self._pipe.send_queue.put(unihash)

	async def get_result(self):
	r = await self._pipe.recv_queue.get()
	return r == "true"


	class GcMarkStream(Stream):
	def __init__(self, pipe, mark):
	super().__init__(pipe)
	self.mark = mark

	async def _send_batch_input(self, i):
	def row_to_dict(row):
	pairs = row.split()
	return dict(zip(pairs[::2], pairs[1::2]))

	await self.send_mark(row_to_dict(i))

	async def send_mark(self, where):
	await self._pipe.send_queue.put(json.dumps({"mark": self.mark, "where": where}))

	async def get_result(self):
	r = await self._pipe.recv_queue.get()
	return json.loads(r)


	class Batch(object):
	def __init__(self, send_queue, recv_queue):
	self.send_queue = send_queue
	self.recv_queue = recv_queue
	self.fill_done = False
	self.send_done = False
	self.cond = asyncio.Condition()
	self.pending = []
	self.sent_count = 0
	self.recv_count = 0
	self.item = None

	async def recv(self, socket):
	while True:
	async with self.cond:
	await self.cond.wait_for(lambda: self.pending or self.send_done)
	if not self.pending:
	if self.send_done:
	return
	continue

	m = await socket.recv()
	await self.recv_queue.put(m)

	async with self.cond:
	self.recv_count += 1
	self.pending.pop(0)

	async def fill(self):
	async for m in self.send_queue:
	async with self.cond:
	# Wait for item to be consumed
	await self.cond.wait_for(lambda: self.item is None)
	self.item = m
	self.cond.notify_all()

	async with self.cond:
	self.fill_done = True
	self.cond.notify_all()

	async def send(self, socket):
	# In the event of a restart due to a reconnect, all in-flight
	# messages need to be resent first to keep to result count in sync
	async with self.cond:
	for m in self.pending:
	await socket.send(m)

	while True:
	async with self.cond:
	await self.cond.wait_for(
	lambda: self.item is not None or self.fill_done
	)
	if self.item is None:
	if self.fill_done:
	self.send_done = True
	self.cond.notify_all()
	return
	continue

	m = self.item

	await socket.send(m)

	async with self.cond:
	self.item = None
	self.pending.append(m)
	self.sent_count += 1
	self.cond.notify_all()

	async def stream(self, socket):
	await bb.asyncrpc.TaskGroup.run(self.send(socket), self.recv(socket))

	def check(self):
	if self.sent_count != self.recv_count:
	raise ConnectionError(
	f"Sent {self.sent_count} messages but only received {self.recv_count}"
	)


	class AsyncClient(bb.asyncrpc.AsyncClient):
	MODE_NORMAL = 0
	MODE_GET_STREAM = 1
	MODE_EXIST_STREAM = 2
	MODE_MARK_STREAM = 3

	def __init__(self, username=None, password=None):
	super().__init__("OEHASHEQUIV", "1.1", logger)
	self.mode = self.MODE_NORMAL
	self.username = username
	self.password = password
	self.saved_become_user = None

	async def setup_connection(self):
	await super().setup_connection()
	self.mode = self.MODE_NORMAL
	if self.username:
	# Save off become user temporarily because auth() resets it
	become = self.saved_become_user
	await self.auth(self.username, self.password)

	if become:
	await self.become_user(become)

	@asynccontextmanager
	async def send_stream(self, mode):
	send_queue = AsyncQueue()
	recv_queue = AsyncQueue()
	b = Batch(send_queue, recv_queue)

	async def proc():
	await self._set_mode(mode)
	await b.stream(self.socket)

	async def process():
	try:
	await self._send_wrapper(proc)
	finally:
	await recv_queue.done()

	# Create background process to process messages
	async with bb.asyncrpc.TaskGroup() as group:
	group.create_task(process())
	group.create_task(b.fill())
	try:
	yield AsyncPipe(send_queue, recv_queue)
	b.check()
	except AsyncQueue.Shutdown as e:
	pass
	finally:
	await send_queue.done()
	await recv_queue.done()

	async def invoke(self, args, skip_mode=False, *kwargs):
	# It's OK if connection errors cause a failure here, because the mode
	# is also reset to normal on a new connection
	if not skip_mode:
	await self._set_mode(self.MODE_NORMAL)
	return await super().invoke(args, *kwargs)

	async def _set_mode(self, new_mode):
	async def stream_to_normal():
	# Check if already in normal mode (e.g. due to a connection reset)
	if self.mode == self.MODE_NORMAL:
	return "ok"
	await self.socket.send("END")
	return await self.socket.recv()

	async def normal_to_stream(command):
	r = await self.invoke({command: None}, skip_mode=True)
	if r != "ok":
	self.check_invoke_error(r)
	raise ConnectionError(
	f"Unable to transition to stream mode: Bad response from server {r!r}"
	)
	self.logger.debug("Mode is now %s", command)

	if new_mode == self.mode:
	return

	self.logger.debug("Transitioning mode %s -> %s", self.mode, new_mode)

	# Always transition to normal mode before switching to any other mode
	if self.mode != self.MODE_NORMAL:
	r = await self._send_wrapper(stream_to_normal)
	if r != "ok":
	self.check_invoke_error(r)
	raise ConnectionError(
	f"Unable to transition to normal mode: Bad response from server {r!r}"
	)
	self.logger.debug("Mode is now normal")

	if new_mode == self.MODE_GET_STREAM:
	await normal_to_stream("get-stream")
	elif new_mode == self.MODE_EXIST_STREAM:
	await normal_to_stream("exists-stream")
	elif new_mode == self.MODE_MARK_STREAM:
	await normal_to_stream("gc-mark-stream")
	elif new_mode != self.MODE_NORMAL:
	raise Exception("Undefined mode transition {self.mode!r} -> {new_mode!r}")

	self.mode = new_mode

	async def get_unihash(self, method, taskhash):
	async with self.get_unihash_stream() as stream:
	await stream.send_query(method, taskhash)
	return await stream.get_result()

	async def get_unihash_batch(self, args):
	async with self.get_unihash_stream() as stream:
	return await stream.batch(args)

	@asynccontextmanager
	async def get_unihash_stream(self):
	async with self.send_stream(self.MODE_GET_STREAM) as pipe:
	yield GetUnihashStream(pipe)

	async def report_unihash(self, taskhash, method, outhash, unihash, extra={}):
	m = extra.copy()
	m["taskhash"] = taskhash
	m["method"] = method
	m["outhash"] = outhash
	m["unihash"] = unihash
	return await self.invoke({"report": m})

	async def report_unihash_equiv(self, taskhash, method, unihash, extra={}):
	m = extra.copy()
	m["taskhash"] = taskhash
	m["method"] = method
	m["unihash"] = unihash
	return await self.invoke({"report-equiv": m})

	async def get_taskhash(self, method, taskhash, all_properties=False):
	return await self.invoke(
	{"get": {"taskhash": taskhash, "method": method, "all": all_properties}}
	)

	async def unihash_exists(self, unihash):
	async with self.unihash_exists_stream() as stream:
	await stream.send_query(unihash)
	return await stream.get_result()

	async def unihash_exists_batch(self, unihashes):
	async with self.unihash_exists_stream() as stream:
	return await stream.batch(unihashes)

	@asynccontextmanager
	async def unihash_exists_stream(self):
	async with self.send_stream(self.MODE_EXIST_STREAM) as pipe:
	yield UnihashExistsStream(pipe)

	async def get_outhash(self, method, outhash, taskhash, with_unihash=True):
	return await self.invoke(
	{
	"get-outhash": {
	"outhash": outhash,
	"taskhash": taskhash,
	"method": method,
	"with_unihash": with_unihash,
	}
	}
	)

	async def get_stats(self):
	return await self.invoke({"get-stats": None})

	async def reset_stats(self):
	return await self.invoke({"reset-stats": None})

	async def backfill_wait(self):
	return (await self.invoke({"backfill-wait": None}))["tasks"]

	async def remove(self, where):
	return await self.invoke({"remove": {"where": where}})

	async def clean_unused(self, max_age):
	return await self.invoke({"clean-unused": {"max_age_seconds": max_age}})

	async def auth(self, username, token):
	result = await self.invoke({"auth": {"username": username, "token": token}})
	self.username = username
	self.password = token
	self.saved_become_user = None
	return result

	async def refresh_token(self, username=None):
	m = {}
	if username:
	m["username"] = username
	result = await self.invoke({"refresh-token": m})
	if (
	self.username
	and not self.saved_become_user
	and result["username"] == self.username
	):
	self.password = result["token"]
	return result

	async def set_user_perms(self, username, permissions):
	return await self.invoke(
	{"set-user-perms": {"username": username, "permissions": permissions}}
	)

	async def get_user(self, username=None):
	m = {}
	if username:
	m["username"] = username
	return await self.invoke({"get-user": m})

	async def get_all_users(self):
	return (await self.invoke({"get-all-users": {}}))["users"]

	async def new_user(self, username, permissions):
	return await self.invoke(
	{"new-user": {"username": username, "permissions": permissions}}
	)

	async def delete_user(self, username):
	return await self.invoke({"delete-user": {"username": username}})

	async def become_user(self, username):
	result = await self.invoke({"become-user": {"username": username}})
	if username == self.username:
	self.saved_become_user = None
	else:
	self.saved_become_user = username
	return result

	async def get_db_usage(self):
	return (await self.invoke({"get-db-usage": {}}))["usage"]

	async def get_db_query_columns(self):
	return (await self.invoke({"get-db-query-columns": {}}))["columns"]

	async def gc_status(self):
	return await self.invoke({"gc-status": {}})

	async def gc_mark(self, mark, where):
	"""
	Starts a new garbage collection operation identified by "mark". If
	garbage collection is already in progress with "mark", the collection
	is continued.

	All unihash entries that match the "where" clause are marked to be
	kept. In addition, any new entries added to the database after this
	command will be automatically marked with "mark"
	"""
	return await self.invoke({"gc-mark": {"mark": mark, "where": where}})

	async def gc_mark_batch(self, mark, rows):
	"""
	Similar to `gc-mark`, but accepts a list of "where" key-value pair
	conditions. It utilizes stream mode to mark hashes, which helps reduce
	the impact of latency when communicating with the hash equivalence
	server.
	"""
	async with self.gc_mark_stream(mark) as stream:
	results = await stream.batch(rows)

	return {"count": sum(int(r["count"]) for r in results)}

	@asynccontextmanager
	async def gc_mark_stream(self, mark):
	async with self.send_stream(self.MODE_MARK_STREAM) as pipe:
	yield GcMarkStream(pipe, mark)

	async def gc_sweep(self, mark):
	"""
	Finishes garbage collection for "mark". All unihash entries that have
	not been marked will be deleted.

	It is recommended to clean unused outhash entries after running this to
	cleanup any dangling outhashes
	"""
	return await self.invoke({"gc-sweep": {"mark": mark}})


	class Client(bb.asyncrpc.Client):
	def __init__(self, username=None, password=None):
	self.username = username
	self.password = password

	super().__init__()
	self._add_methods(
	"connect_tcp",
	"connect_websocket",
	"get_unihash",
	"get_unihash_batch",
	"report_unihash",
	"report_unihash_equiv",
	"get_taskhash",
	"unihash_exists",
	"unihash_exists_batch",
	"get_outhash",
	"get_stats",
	"reset_stats",
	"backfill_wait",
	"remove",
	"clean_unused",
	"auth",
	"refresh_token",
	"set_user_perms",
	"get_user",
	"get_all_users",
	"new_user",
	"delete_user",
	"become_user",
	"get_db_usage",
	"get_db_query_columns",
	"gc_status",
	"gc_mark",
	"gc_mark_batch",
	"gc_sweep",
	)

	def _get_async_client(self):
	return AsyncClient(self.username, self.password)