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
def send_pch_file(fileobj): send_file(self.sender.send_msg, fileobj)