def got_data_from_server(self, msg): if msg[0] == b'SESSION_CANCELLED': assert self.cancelled self.__complete(SessionResult.cancelled) return True elif msg[0] == b'TIMED_OUT': self.__complete(SessionResult.timed_out) return True # It is possible that cancellation arrived too late, # that the server already sent the final message and # unregistered its session. In that case we will never # get confirmation. # This state requires a response, so the session must be still alive # on the server. if self.state == self.STATE_WAIT_FOR_MISSING_FILES: assert len(msg) == 3 and msg[1] == b'MISSING_FILES' assert self.sender is None self.sender = self.Sender(self.send_msg, msg[0].tobytes()) if self.cancelled: self.sender.send_msg([b'CANCEL_SESSION']) missing_files, need_compiler, need_pch = pickle.loads(msg[2].memory()) task_files = [b'TASK_FILES'] task_files.extend(self.task_files_bundle(missing_files)) self.sender.send_msg(task_files) if need_compiler: zip_data = BytesIO() with zipfile.ZipFile(zip_data, mode='w') as zip_file: for path, file in self.task.compiler_info.files: if path: zip_file.write(path.decode(), file.decode()) send_file(self.sender.send_msg, BytesIO(zip_data.getbuffer())) del zip_data if need_pch: assert self.task.pch_file is not None def send_pch_file(fileobj): send_file(self.sender.send_msg, fileobj) pch_file = os.path.join(os.getcwd(), self.task.pch_file[0]) self.compressor.compress_file(pch_file, send_pch_file) self.state = self.STATE_WAIT_FOR_SERVER_RESPONSE elif self.state == self.STATE_WAIT_FOR_SERVER_RESPONSE: server_status = msg[0] if server_status == b'SERVER_FAILED': self.retcode = -1 self.stdout = b'' self.stderr = msg[1].tobytes() self.__complete(SessionResult.failure) return True else: assert server_status == b'SERVER_DONE' self.retcode, self.stdout, self.stderr, server_times = pickle.loads(msg[1].memory()) logging.debug("Got {} retcode".format(self.retcode)) for name, duration in server_times.items(): self.timer.add_time(name, duration) if self.task.register_completion(self): assert not self.cancelled if self.retcode == 0: self.sender.send_msg([b'SEND_CONFIRMATION', b'\x01']) self.obj_desc = BytesIO() self.state = self.STATE_RECEIVE_OBJECT_FILE self.result_futures = [] self.receive_result_time = SimpleTimer() else: self.__complete(SessionResult.success) return True else: if self.retcode == 0: self.sender.send_msg([b'SEND_CONFIRMATION', b'\x00']) self.__complete(SessionResult.too_late) return True elif self.state == self.STATE_RECEIVE_OBJECT_FILE: assert not self.cancelled more, data = msg self.obj_desc.write(data.memory()) if more == b'\x00': file_index = len(self.result_futures) future = self.loop.run_in_executor(self.executor, self._decompress_to_disk, self.obj_desc, self.task.result_files[file_index]) self.result_futures.append(future) # Complete the session once all files are decompressed. if len(self.result_futures) == len(self.task.result_files): self.timer.add_time('download result files', self.receive_result_time.get()) asyncio.gather(*self.result_futures, loop=self.loop ).add_done_callback(self.complete_session) return True else: # Prepare to receive another file. self.obj_desc = BytesIO() else: assert not "Invalid state" return False
class ServerSession: STATE_START = 0 STATE_WAIT_FOR_MISSING_FILES = 1 STATE_RECEIVE_OBJECT_FILE = 2 STATE_WAIT_FOR_SERVER_RESPONSE = 3 STATE_FINISH = 4 class Sender: def __init__(self, send_msg, session_id): self._session_id = session_id self._send_msg = send_msg def send_msg(self, data): self._send_msg([self._session_id] + list(data)) def __init__(self, session_id, task, send_msg, node, loop, executor, compressor, completion_callback): self.state = self.STATE_START self.task = task self.node = node self.task.register_session(self) self.cancelled = False self.loop = loop self.executor = executor self.compressor = compressor self.result = None self.local_id = session_id self.send_msg = send_msg self.sender = None self.completion_callback = completion_callback def start(self): assert self.state == self.STATE_START self.send_msg([b'NEW_SESSION', self.local_id, b'SERVER_TASK', pickle.dumps(self.task.server_task)]) self.state = self.STATE_WAIT_FOR_MISSING_FILES self.time_started = time() def cancel(self): if self.sender: self.sender.send_msg([b'CANCEL_SESSION']) self.cancelled = True def terminate(self): self.state = self.STATE_FINISH self.__complete(SessionResult.terminated) def __complete(self, result): self.state = self.STATE_FINISH self.time_completed = time() self.result = result self.completion_callback(self) @property def timer(self): return self.node.timer() def got_data_from_server(self, msg): if msg[0] == b'SESSION_CANCELLED': assert self.cancelled self.__complete(SessionResult.cancelled) return True elif msg[0] == b'TIMED_OUT': self.__complete(SessionResult.timed_out) return True # It is possible that cancellation arrived too late, # that the server already sent the final message and # unregistered its session. In that case we will never # get confirmation. # This state requires a response, so the session must be still alive # on the server. if self.state == self.STATE_WAIT_FOR_MISSING_FILES: assert len(msg) == 3 and msg[1] == b'MISSING_FILES' assert self.sender is None self.sender = self.Sender(self.send_msg, msg[0].tobytes()) if self.cancelled: self.sender.send_msg([b'CANCEL_SESSION']) missing_files, need_compiler, need_pch = pickle.loads(msg[2].memory()) task_files = [b'TASK_FILES'] task_files.extend(self.task_files_bundle(missing_files)) self.sender.send_msg(task_files) if need_compiler: zip_data = BytesIO() with zipfile.ZipFile(zip_data, mode='w') as zip_file: for path, file in self.task.compiler_info.files: if path: zip_file.write(path.decode(), file.decode()) send_file(self.sender.send_msg, BytesIO(zip_data.getbuffer())) del zip_data if need_pch: assert self.task.pch_file is not None def send_pch_file(fileobj): send_file(self.sender.send_msg, fileobj) pch_file = os.path.join(os.getcwd(), self.task.pch_file[0]) self.compressor.compress_file(pch_file, send_pch_file) self.state = self.STATE_WAIT_FOR_SERVER_RESPONSE elif self.state == self.STATE_WAIT_FOR_SERVER_RESPONSE: server_status = msg[0] if server_status == b'SERVER_FAILED': self.retcode = -1 self.stdout = b'' self.stderr = msg[1].tobytes() self.__complete(SessionResult.failure) return True else: assert server_status == b'SERVER_DONE' self.retcode, self.stdout, self.stderr, server_times = pickle.loads(msg[1].memory()) logging.debug("Got {} retcode".format(self.retcode)) for name, duration in server_times.items(): self.timer.add_time(name, duration) if self.task.register_completion(self): assert not self.cancelled if self.retcode == 0: self.sender.send_msg([b'SEND_CONFIRMATION', b'\x01']) self.obj_desc = BytesIO() self.state = self.STATE_RECEIVE_OBJECT_FILE self.result_futures = [] self.receive_result_time = SimpleTimer() else: self.__complete(SessionResult.success) return True else: if self.retcode == 0: self.sender.send_msg([b'SEND_CONFIRMATION', b'\x00']) self.__complete(SessionResult.too_late) return True elif self.state == self.STATE_RECEIVE_OBJECT_FILE: assert not self.cancelled more, data = msg self.obj_desc.write(data.memory()) if more == b'\x00': file_index = len(self.result_futures) future = self.loop.run_in_executor(self.executor, self._decompress_to_disk, self.obj_desc, self.task.result_files[file_index]) self.result_futures.append(future) # Complete the session once all files are decompressed. if len(self.result_futures) == len(self.task.result_files): self.timer.add_time('download result files', self.receive_result_time.get()) asyncio.gather(*self.result_futures, loop=self.loop ).add_done_callback(self.complete_session) return True else: # Prepare to receive another file. self.obj_desc = BytesIO() else: assert not "Invalid state" return False @classmethod def _decompress_to_disk(cls, fileobj, output): fileobj.seek(0) decompressor = zlib.decompressobj() with open(output, "wb") as obj: for data in iter(lambda : fileobj.read(256 * 1024), b''): obj.write(decompressor.decompress(data)) obj.write(decompressor.flush()) logging.debug("Written {} bytes".format(obj.tell())) def complete_session(self, future): try: # We might have been terminated. if self.state == self.STATE_FINISH: assert self.result == SessionResult.terminated return future_count = len(future.result()) except Exception as e: logging.exception(e) self.__complete(SessionResult.failure) else: assert future_count == len(self.task.result_files) self.__complete(SessionResult.success) def task_files_bundle(self, in_filelist): header_info = self.task.header_info source_file = self.task.source result = [] for dir, data in header_info: dir_bytes = dir.encode() for file, relative, content_entry in data: if not relative and not (dir, file) in in_filelist: continue result.extend((dir_bytes, file.encode(), content_entry.buffer())) with open(source_file, 'rb') as src: result.extend((b'', source_file.encode(), src.read())) return result def get_info(self): assert self.state == self.STATE_FINISH assert self.result is not None return { "hostname" : self.node.node_dict()['hostname'], "port" : self.node.node_dict()['port'], "started" : self.time_started, "completed" : self.time_completed, "result" : self.result, }