async def read(self, call_id: int = 0) -> DataWrapper[bytes]:
if call_id > 0:
data = self._call_store.pop(call_id)
if data is not None:
return data
try:
self.warning("+++++++++ READ MUTEX LOCK WAIT")
async with self._read_lock:
self.warning("++++++ READ")
hsize, bsize = await self._read_size()
if hsize == 0:
return DataWrapper(
False, "", b'', 104,
"Connection closed by peer (got 0 size)")
data = await self._read(hsize + bsize)
msg = transport_pb2.TransportHeader()
msg.ParseFromString(data[:hsize])
if msg.status.success:
response = DataWrapper(True, msg.uri, data[hsize:], id=msg.id)
else:
response = DataWrapper(False,
msg.uri,
b'',
msg.status.error_code,
"",
msg.status.narrative[:],
id=msg.id)
if call_id != 0 and msg.id != call_id:
self._call_store.put(call_id, response)
return await self.read(call_id)
return response
except ConnectionResetError as e:
return DataWrapper(False, "", b'', 104, str(e))
async def handle_message(
self, msg: remove_pb2.Remove
) -> DataWrapper[response_pb2.UpdateResponse]:
global _count_success, _count_request, _count_exception, _count_fails
_count_request.inc()
resp = response_pb2.RemoveResponse()
try:
if msg.all:
status = self.dap_manager.removeRow(msg.key, msg.agent_key)
self.info("All data removed for: tcp:///{}/{}".format(
msg.key, msg.agent_key))
else:
upd = self.proto_wrapper.get_instance(toUpdate(msg))
status = self.dap_manager.remove(upd.toDapUpdate())
dms = ", ".join(map(lambda x: x.model.name, msg.data_models))
self.info(
"Data removed for: tcp:///{}/{}, data models=[{}]".format(
msg.key, msg.agent_key, dms))
if status:
resp.status = ResponseType.Value("SUCCESS")
_count_success.inc()
else:
resp.status = ResponseType.Value("NOT_FOUND")
_count_fails.inc()
except Exception as e:
self.error("Remove failed because: ", str(e))
resp.status = ResponseType.Value("ERROR")
resp.message = str(e)
_count_exception.inc()
return DataWrapper(True, "remove", resp)
def process_update(self, msg) -> DataWrapper[response_pb2.UpdateResponse]:
global _count_success, _count_request, _count_exception, _count_fails
_count_request.inc()
response = DataWrapper(False, "update", response_pb2.UpdateResponse())
try:
upd = self.proto_wrapper.get_instance(msg)
status = self.dap_manager.update(upd.toDapUpdate())
response.success = status
if status:
response.data.status = ResponseType.Value("SUCCESS")
_count_success.inc()
else:
response.data.status = ResponseType.Value("ERROR")
_count_fails.inc()
except InvalidAttribute as e:
msg = "Got invalid attribute: " + str(e)
response.data.status = ResponseType.Value("INVALID_ATTRIBUTE")
response.data.message = msg
self.log.info(msg)
_count_exception.inc()
except MissingAddress as e:
response.data.status = ResponseType.Value("MISSING_ADDRESS")
msg = str(e)
response.data.message = msg
self.log.info(msg)
_count_exception.inc()
except Exception as e:
msg = "Failed to update data, because: " + str(e)
response.data.status = ResponseType.Value("ERROR")
response.data.message = msg
self.exception(e)
_count_exception.inc()
return response
async def handle_message(
self,
msg: core_pb2.CoreLocation) -> DataWrapper[core_pb2.Response]:
response = DataWrapper(False, "location", core_pb2.Response())
update = dap_update_pb2.DapUpdate()
try:
upd = update.update.add()
upd.tablename = self._db_structure["table"]
upd.fieldname = self._db_structure["field"]
upd.key.core = msg.core_key
upd.value.type = 9
upd.value.l.CopyFrom(msg.location)
except Exception as e:
msg = "Failed to create dap_update proto from director location update: ", str(
e)
self.warning(msg)
response.error_code = 500
response.add_narrative(msg)
try:
status = self._dap_manager.update(update)
response.success = status
except Exception as e:
msg = "DapManager: failed update: " + str(e)
response.add_narrative(msg)
self.warning(msg)
return response
async def handle_message(self, msg) -> DataWrapper[Any]:
response = await self._node.broadcast(self._path, msg)
if type(response) != list:
response = [response]
response = self.__flatten_list(response)
results = await asyncio.gather(*[
self._serializer.deserialize(serialized) for serialized in response
])
return [DataWrapper(True, "", x) for x in results]
async def handle_message(self, msg: node_pb2.PeerUpdate) -> DataWrapper[core_pb2.Response]:
response = DataWrapper(True, "peer", core_pb2.Response())
try:
self.info("ADD PEERS: ", msg)
for peer in msg.add_peers:
self._manager.add_peer(peer.name, peer.host, peer.port)
except Exception as e:
self.exception(e)
response.success = False
return response
async def handle_message(self, msg):
response = await self._node.broadcast(self._path, copy.deepcopy(msg))
if type(response) != list:
response = [response]
response = self.__flatten_list(response)
transformed = []
for r in response:
if len(r.data) < 1:
deserialized = None
else:
deserialized = await self._serializer.deserialize(r.data)
transformed.append(
DataWrapper(r.success, r.uri, deserialized, r.error_code, "",
r.narrative, r.id))
return transformed
async def handle_message(
self, msg: info_pb2.NodeInfoRequest
) -> DataWrapper[info_pb2.NodeInfoResponse]:
response = DataWrapper(True, "info", info_pb2.NodeInfoResponse())
try:
self.info("GOT INFO request! ")
response.data.search_key = self._node_config["search_key"]
response.data.core_key = self._node_config["core_key"]
response.data.host = self._node_config["host"]
response.data.core_port = self._node_config["core_port"]
response.data.search_port = self._node_config["search_port"]
except Exception as e:
self.exception(e)
response.success = False
return response
async def handle_message(
self, msg: remove_pb2.Remove
) -> DataWrapper[response_pb2.UpdateResponse]:
resp = response_pb2.RemoveResponse()
try:
if msg.all:
status = self.dap_manager.removeAll(msg.key)
else:
upd = self.proto_wrapper.get_instance(toUpdate(msg))
status = self.dap_manager.remove(upd.toDapUpdate())
if status:
resp.status = ResponseType.Value("SUCCESS")
else:
resp.status = ResponseType.Value("NOT_FOUND")
except Exception as e:
resp.status = ResponseType.Value("ERROR")
resp.message = str(e)
return DataWrapper(True, "remove", resp)
def merge_response(self, resps: List[DataWrapper[response_pb2.SearchResponse]]) -> \
DataWrapper[response_pb2.SearchResponse]:
resp = response_pb2.SearchResponse()
error_code = 1000
narrative = []
keys = []
success = False
for r in resps:
result = []
success |= r.success
error_code = min(error_code, r.error_code)
narrative.extend(r.narrative)
for rr in r.data.result:
if rr.key in keys:
continue
keys.append(rr.key)
result.append(rr)
resp.result.extend(result)
return DataWrapper(success, "search", resp, error_code, "", narrative)
async def handle_message(self, msg: Any) -> DataWrapper[Any]:
response = DataWrapper(False, "", None)
try:
if isinstance(msg, dap_interface_pb2.NoInputParameter):
proto = self._dap_call()
else:
proto = self._dap_call(msg)
response.success = True
response.data = proto
except Exception as e:
error = dap_interface_pb2.Successfulness()
error.success = False
error.errorcode = 503
error.narrative.append(str(e))
self.exception("Exception during DAP call: ", e)
response.error_code = 503
response.add_narrative(str(e))
response.data = error
return response
async def route(self, path: str, data) -> DataWrapper[bytes]:
if path in self.__routing_serializer:
serializer = self.__routing_serializer[path]
msg = await serializer.deserialize(data)
if path in self.__routing_handler:
cos = []
for handler in self.__routing_handler[path]:
cos.append(asyncio.create_task(
handler.handle_message(msg)))
response_list = await asyncio.gather(*cos,
return_exceptions=True)
if len(response_list) == 1:
response = response_list[0]
else:
protos_by_type = {}
for p in response_list:
if p is None:
continue
if isinstance(p, Exception):
self.log.exception(
"Exception happened in handler for path {}: {}"
.format(path, str(p)))
continue
elif isinstance(p, list):
if len(p) == 0:
continue
d = self.__flatten_list(p)
else:
d = [p]
for e in d:
if not e.success or e.data is None:
self.info("Error in handlers: code: ",
e.error_code, ", message: ", e.msg(),
", data: ", e.data)
continue
protos_by_type.setdefault(e.data.__class__,
[]).append(e)
merged_list = []
for k in protos_by_type:
if len(protos_by_type[k]) == 0:
continue
if k in self.__response_merger:
merged_list.append(
self.__response_merger[k].merge_response(
protos_by_type[k]))
else:
for e in protos_by_type[k]:
merged_list.append(e)
if len(merged_list) == 1:
response = merged_list[0]
elif len(merged_list) == 0:
self.log.warning("Empty merged list")
return DataWrapper(False, path, b'', 61,
"Merged list is empty")
else:
self.info("Using response builder for path ", path,
" for list: ", [p.data for p in merged_list])
response = self.__response_builder[
path].build_responses(merged_list)
if not response.success:
self.warning(
"Error in route handler for path %s, error_code %d, message: %s",
path, response.error_code, response.msg())
wrapped_data = Serializable(response.data)
if isinstance(data, dict):
wrapped_data.target_type = Serializable.TargetType.JSON
serialized_response = await serializer.serialize(wrapped_data)
return DataWrapper(response.success,
path,
serialized_response,
response.error_code,
"",
response.narrative,
id=response.id)
else:
msg = "Message handler not register for path: %s" % path
self.log.error(msg)
return DataWrapper(False, path, b'', 53, msg)
else:
msg = "Serializer not registered for path: %s" % path
self.log.error(msg)
return DataWrapper(False, path, b'', 53, msg)
async def handle_message(
self,
msg: query_pb2.Query) -> DataWrapper[response_pb2.SearchResponse]:
global _count_request, _count_success, _count_exception, _count_ignored, _count_empty
resp = response_pb2.SearchResponse()
distance = None
_count_request.inc()
try:
location = self._dap_manager.getPlaneInformation(
"location")["values"]
if len(location) > 1:
self.warning(
"Got more then 1 location from dapManager (I'm using first, ignoring the rest now): ",
location)
#TODO: multiple core
location = location[0].value.l
prev_distance = msg.directed_search.distance.geo
target = msg.directed_search.target.geo
distance = geo_distance(location, target)
if prev_distance > 0 and distance > prev_distance:
#ignore query
_count_ignored.inc()
self.warning(
"Ignoring query with TTL %d, because source (%s) distance (%.3f) was smaller then "
"my distance (%.3f)", msg.ttl,
msg.source_key.decode("UTF-8"), prev_distance, distance)
return DataWrapper(False, "search", resp, 601,
"Ignoring query because distance!")
else:
self.info(
"Handling query with TTL %d, because source (%s) distance (%.3f) was greater then"
" my distance (%.3f)", msg.ttl,
msg.source_key.decode("UTF-8"), prev_distance, distance)
except Exception as e:
#TODO: what?
self.info("Exception during location check: ", str(e))
#return DataWrapper(False, "search", resp, 600, "Ignoring query, location check error!")
_count_exception.inc()
try:
query = self._proto_wrapper.get_instance(msg)
result = self._dap_manager.execute(query.toDapQuery())
items = {}
for element in result.identifiers:
core, agent_id = element.core, element.agent
if core not in items:
item = response_pb2.SearchResponse.Item()
address = element.uri.split(":")
if len(address) == 2:
item.ip = address[0]
item.port = int(address[1])
elif len(address) == 3:
item.ip = address[1]
item.port = int(address[2])
else:
self.warning("No valid address for core: ", core,
"! address=", address)
continue
item.key = core
items[core] = item
if distance:
#TODO multiple core
item.distance = distance
agent = items[core].agents.add()
agent.key = agent_id
agent.score = element.score
resp.result.extend(items.values())
self.info("SEARCH RESULT: ", resp)
if len(items) == 0:
_count_empty.inc()
except Exception as e:
self.exception("Exception while processing query: ", str(e))
_count_exception.inc()
return DataWrapper(False, "search", resp, 500, str(e))
_count_success.inc()
return DataWrapper(True, "search", resp)