class NailgunClientSession(NailgunProtocol): """Handles a single nailgun client session.""" def __init__(self, sock, in_fd, out_fd, err_fd): self._sock = sock self._input_reader = NailgunStreamReader(in_fd, self._sock) if in_fd else None self._stdout = out_fd self._stderr = err_fd self.remote_pid = None def _maybe_start_input_reader(self): if self._input_reader: self._input_reader.start() def _maybe_stop_input_reader(self): if self._input_reader: self._input_reader.stop() def _process_session(self): """Process the outputs of the nailgun session.""" try: for chunk_type, payload in self.iter_chunks(self._sock, return_bytes=True): if chunk_type == ChunkType.STDOUT: self._stdout.write(payload) self._stdout.flush() elif chunk_type == ChunkType.STDERR: self._stderr.write(payload) self._stderr.flush() elif chunk_type == ChunkType.EXIT: self._stdout.flush() self._stderr.flush() return int(payload) elif chunk_type == ChunkType.PID: self.remote_pid = int(payload) elif chunk_type == ChunkType.START_READING_INPUT: self._maybe_start_input_reader() else: raise self.ProtocolError( 'received unexpected chunk {} -> {}'.format( chunk_type, payload)) finally: # Bad chunk types received from the server can throw NailgunProtocol.ProtocolError in # NailgunProtocol.iter_chunks(). This ensures the NailgunStreamReader is always stopped. self._maybe_stop_input_reader() def execute(self, working_dir, main_class, *arguments, **environment): # Send the nailgun request. self.send_request(self._sock, working_dir, main_class, *arguments, **environment) # Process the remainder of the nailgun session. return self._process_session()
class NailgunClientSession(NailgunProtocol): """Handles a single nailgun client session.""" def __init__(self, sock, in_fd, out_fd, err_fd): self._sock = sock self._input_reader = NailgunStreamReader(in_fd, self._sock) if in_fd else None self._stdout = out_fd self._stderr = err_fd self.remote_pid = None def _maybe_start_input_reader(self): if self._input_reader: self._input_reader.start() def _maybe_stop_input_reader(self): if self._input_reader: self._input_reader.stop() def _process_session(self): """Process the outputs of the nailgun session.""" try: for chunk_type, payload in self.iter_chunks(self._sock): if chunk_type == ChunkType.STDOUT: self._stdout.write(payload) self._stdout.flush() elif chunk_type == ChunkType.STDERR: self._stderr.write(payload) self._stderr.flush() elif chunk_type == ChunkType.EXIT: self._stdout.flush() self._stderr.flush() return int(payload) elif chunk_type == ChunkType.PID: self.remote_pid = int(payload) elif chunk_type == ChunkType.START_READING_INPUT: self._maybe_start_input_reader() else: raise self.ProtocolError('received unexpected chunk {} -> {}'.format(chunk_type, payload)) finally: # Bad chunk types received from the server can throw NailgunProtocol.ProtocolError in # NailgunProtocol.iter_chunks(). This ensures the NailgunStreamReader is always stopped. self._maybe_stop_input_reader() def execute(self, working_dir, main_class, *arguments, **environment): # Send the nailgun request. self.send_request(self._sock, working_dir, main_class, *arguments, **environment) # Process the remainder of the nailgun session. return self._process_session()
class NailgunClientSession(NailgunProtocol): """Handles a single nailgun client session.""" BUF_SIZE = 8192 def __init__(self, sock, in_fd, out_fd, err_fd): self._sock = sock self._input_reader = NailgunStreamReader(in_fd, self._sock, self.BUF_SIZE) if in_fd else None self._stdout = out_fd self._stderr = err_fd @contextmanager def _maybe_input_reader_running(self): if self._input_reader: self._input_reader.start() yield if self._input_reader: self._input_reader.stop() def _process_session(self): """Process the outputs of the nailgun session.""" for chunk_type, payload in self.iter_chunks(self._sock): if chunk_type == ChunkType.STDOUT: self._stdout.write(payload) self._stdout.flush() elif chunk_type == ChunkType.STDERR: self._stderr.write(payload) self._stderr.flush() elif chunk_type == ChunkType.EXIT: self._stdout.flush() self._stderr.flush() return int(payload) else: raise self.ProtocolError('Received unexpected chunk {} -> {}'.format(chunk_type, payload)) def execute(self, working_dir, main_class, *arguments, **environment): # Send the nailgun request. self.send_request(self._sock, working_dir, main_class, *arguments, **environment) # Launch the NailgunStreamReader if applicable and process the remainder of the nailgun session. with self._maybe_input_reader_running(): return self._process_session()
class NailgunClientSession(NailgunProtocol): """Handles a single nailgun client session.""" def __init__(self, sock, in_fd, out_fd, err_fd, exit_on_broken_pipe=False): self._sock = sock self._input_reader = NailgunStreamReader(in_fd, self._sock) if in_fd else None self._stdout = out_fd self._stderr = err_fd self._exit_on_broken_pipe = exit_on_broken_pipe self.remote_pid = None def _maybe_start_input_reader(self): if self._input_reader: self._input_reader.start() def _maybe_stop_input_reader(self): if self._input_reader: self._input_reader.stop() def _write_flush(self, fd, payload=None): """Write a payload to a given fd (if provided) and flush the fd.""" try: if payload: fd.write(payload) fd.flush() except (IOError, OSError) as e: # If a `Broken Pipe` is encountered during a stdio fd write, we're headless - bail. if e.errno == errno.EPIPE and self._exit_on_broken_pipe: sys.exit() # Otherwise, re-raise. raise def _process_session(self): """Process the outputs of the nailgun session.""" try: for chunk_type, payload in self.iter_chunks(self._sock, return_bytes=True): if chunk_type == ChunkType.STDOUT: self._write_flush(self._stdout, payload) elif chunk_type == ChunkType.STDERR: self._write_flush(self._stderr, payload) elif chunk_type == ChunkType.EXIT: self._write_flush(self._stdout) self._write_flush(self._stderr) return int(payload) elif chunk_type == ChunkType.PID: self.remote_pid = int(payload) elif chunk_type == ChunkType.START_READING_INPUT: self._maybe_start_input_reader() else: raise self.ProtocolError('received unexpected chunk {} -> {}'.format(chunk_type, payload)) finally: # Bad chunk types received from the server can throw NailgunProtocol.ProtocolError in # NailgunProtocol.iter_chunks(). This ensures the NailgunStreamReader is always stopped. self._maybe_stop_input_reader() def execute(self, working_dir, main_class, *arguments, **environment): # Send the nailgun request. self.send_request(self._sock, working_dir, main_class, *arguments, **environment) # Process the remainder of the nailgun session. return self._process_session()