blob: 1a67c6982d3f817fb90adc9c921da7ef97bd7512 [file] [log] [blame]
Brad Bishopa34c0302019-09-23 22:34:48 -04001# Copyright (C) 2019 Garmin Ltd.
2#
3# SPDX-License-Identifier: GPL-2.0-only
4#
5
Brad Bishopa34c0302019-09-23 22:34:48 -04006import logging
7import socket
Andrew Geisslerc926e172021-05-07 16:11:35 -05008import bb.asyncrpc
9from . import create_async_client
Brad Bishopa34c0302019-09-23 22:34:48 -040010
11
Andrew Geissler6ce62a22020-11-30 19:58:47 -060012logger = logging.getLogger("hashserv.client")
Brad Bishopa34c0302019-09-23 22:34:48 -040013
14
Andrew Geisslerc926e172021-05-07 16:11:35 -050015class AsyncClient(bb.asyncrpc.AsyncClient):
Brad Bishopa34c0302019-09-23 22:34:48 -040016 MODE_NORMAL = 0
17 MODE_GET_STREAM = 1
18
19 def __init__(self):
Andrew Geisslerc926e172021-05-07 16:11:35 -050020 super().__init__('OEHASHEQUIV', '1.1', logger)
Brad Bishopa34c0302019-09-23 22:34:48 -040021 self.mode = self.MODE_NORMAL
22
Andrew Geisslerc926e172021-05-07 16:11:35 -050023 async def setup_connection(self):
24 await super().setup_connection()
25 cur_mode = self.mode
26 self.mode = self.MODE_NORMAL
27 await self._set_mode(cur_mode)
Brad Bishopa34c0302019-09-23 22:34:48 -040028
Andrew Geissler6ce62a22020-11-30 19:58:47 -060029 async def send_stream(self, msg):
30 async def proc():
31 self.writer.write(("%s\n" % msg).encode("utf-8"))
32 await self.writer.drain()
33 l = await self.reader.readline()
Brad Bishopa34c0302019-09-23 22:34:48 -040034 if not l:
Andrew Geisslerc926e172021-05-07 16:11:35 -050035 raise ConnectionError("Connection closed")
Andrew Geissler6ce62a22020-11-30 19:58:47 -060036 return l.decode("utf-8").rstrip()
Brad Bishopa34c0302019-09-23 22:34:48 -040037
Andrew Geissler6ce62a22020-11-30 19:58:47 -060038 return await self._send_wrapper(proc)
Brad Bishopa34c0302019-09-23 22:34:48 -040039
Andrew Geissler6ce62a22020-11-30 19:58:47 -060040 async def _set_mode(self, new_mode):
Brad Bishopa34c0302019-09-23 22:34:48 -040041 if new_mode == self.MODE_NORMAL and self.mode == self.MODE_GET_STREAM:
Andrew Geissler6ce62a22020-11-30 19:58:47 -060042 r = await self.send_stream("END")
43 if r != "ok":
Andrew Geisslerc926e172021-05-07 16:11:35 -050044 raise ConnectionError("Bad response from server %r" % r)
Brad Bishopa34c0302019-09-23 22:34:48 -040045 elif new_mode == self.MODE_GET_STREAM and self.mode == self.MODE_NORMAL:
Andrew Geissler6ce62a22020-11-30 19:58:47 -060046 r = await self.send_message({"get-stream": None})
47 if r != "ok":
Andrew Geisslerc926e172021-05-07 16:11:35 -050048 raise ConnectionError("Bad response from server %r" % r)
Brad Bishopa34c0302019-09-23 22:34:48 -040049 elif new_mode != self.mode:
Andrew Geissler6ce62a22020-11-30 19:58:47 -060050 raise Exception(
51 "Undefined mode transition %r -> %r" % (self.mode, new_mode)
52 )
Brad Bishopa34c0302019-09-23 22:34:48 -040053
54 self.mode = new_mode
55
Andrew Geissler6ce62a22020-11-30 19:58:47 -060056 async def get_unihash(self, method, taskhash):
57 await self._set_mode(self.MODE_GET_STREAM)
58 r = await self.send_stream("%s %s" % (method, taskhash))
Brad Bishopa34c0302019-09-23 22:34:48 -040059 if not r:
60 return None
61 return r
62
Andrew Geissler6ce62a22020-11-30 19:58:47 -060063 async def report_unihash(self, taskhash, method, outhash, unihash, extra={}):
64 await self._set_mode(self.MODE_NORMAL)
Brad Bishopa34c0302019-09-23 22:34:48 -040065 m = extra.copy()
Andrew Geissler6ce62a22020-11-30 19:58:47 -060066 m["taskhash"] = taskhash
67 m["method"] = method
68 m["outhash"] = outhash
69 m["unihash"] = unihash
70 return await self.send_message({"report": m})
Brad Bishopa34c0302019-09-23 22:34:48 -040071
Andrew Geissler6ce62a22020-11-30 19:58:47 -060072 async def report_unihash_equiv(self, taskhash, method, unihash, extra={}):
73 await self._set_mode(self.MODE_NORMAL)
Andrew Geissler82c905d2020-04-13 13:39:40 -050074 m = extra.copy()
Andrew Geissler6ce62a22020-11-30 19:58:47 -060075 m["taskhash"] = taskhash
76 m["method"] = method
77 m["unihash"] = unihash
78 return await self.send_message({"report-equiv": m})
Andrew Geissler82c905d2020-04-13 13:39:40 -050079
Andrew Geissler6ce62a22020-11-30 19:58:47 -060080 async def get_taskhash(self, method, taskhash, all_properties=False):
81 await self._set_mode(self.MODE_NORMAL)
82 return await self.send_message(
83 {"get": {"taskhash": taskhash, "method": method, "all": all_properties}}
84 )
Andrew Geissler475cb722020-07-10 16:00:51 -050085
Andrew Geisslerd1e89492021-02-12 15:35:20 -060086 async def get_outhash(self, method, outhash, taskhash):
87 await self._set_mode(self.MODE_NORMAL)
88 return await self.send_message(
89 {"get-outhash": {"outhash": outhash, "taskhash": taskhash, "method": method}}
90 )
91
Andrew Geissler6ce62a22020-11-30 19:58:47 -060092 async def get_stats(self):
93 await self._set_mode(self.MODE_NORMAL)
94 return await self.send_message({"get-stats": None})
Brad Bishopa34c0302019-09-23 22:34:48 -040095
Andrew Geissler6ce62a22020-11-30 19:58:47 -060096 async def reset_stats(self):
97 await self._set_mode(self.MODE_NORMAL)
98 return await self.send_message({"reset-stats": None})
99
100 async def backfill_wait(self):
101 await self._set_mode(self.MODE_NORMAL)
102 return (await self.send_message({"backfill-wait": None}))["tasks"]
103
104
Andrew Geisslerc926e172021-05-07 16:11:35 -0500105class Client(bb.asyncrpc.Client):
Andrew Geissler6ce62a22020-11-30 19:58:47 -0600106 def __init__(self):
Andrew Geisslerc926e172021-05-07 16:11:35 -0500107 super().__init__()
108 self._add_methods(
Andrew Geissler6ce62a22020-11-30 19:58:47 -0600109 "connect_tcp",
Andrew Geissler6ce62a22020-11-30 19:58:47 -0600110 "close",
111 "get_unihash",
112 "report_unihash",
113 "report_unihash_equiv",
114 "get_taskhash",
115 "get_stats",
116 "reset_stats",
117 "backfill_wait",
Andrew Geisslerc926e172021-05-07 16:11:35 -0500118 )
Andrew Geissler6ce62a22020-11-30 19:58:47 -0600119
Andrew Geisslerc926e172021-05-07 16:11:35 -0500120 def _get_async_client(self):
121 return AsyncClient()