Python nailgun_protocol.NailgunProtocol类代码示例

  def run(self):
    while self._in_fds and not self.is_stopped:
      readable, _, errored = select.select(self._in_fds, [], self._in_fds, self._select_timeout)

      if readable:
        for fileno in readable:
          data = os.read(fileno, self._buf_size)

          if not data:
            # We've reached EOF.
            try:
              if self._chunk_eof_type is not None:
                NailgunProtocol.write_chunk(self._socket, self._chunk_eof_type)
            finally:
              try:
                os.close(fileno)
              finally:
                self._in_fds.remove(fileno)
          else:
            NailgunProtocol.write_chunk(
              self._socket,
              self._fileno_chunk_type_map[fileno],
              data
            )

      if errored:
        for fileno in errored:
          self._in_fds.remove(fileno)

 def test_send_exit_default(self):
   NailgunProtocol.send_exit(self.server_sock)
   chunk_type, payload = NailgunProtocol.read_chunk(self.client_sock)
   self.assertEqual(
     (chunk_type, payload),
     (ChunkType.EXIT, self.EMPTY_PAYLOAD)
   )

 def open(cls, sock, isatty=False):
   with _pipe(isatty) as (read_fd, write_fd):
     reader = NailgunStreamStdinReader(sock, os.fdopen(write_fd, 'wb'))
     with reader.running():
       # Instruct the thin client to begin reading and sending stdin.
       NailgunProtocol.send_start_reading_input(sock)
       yield read_fd

 def test_send_start_reading_input(self):
   NailgunProtocol.send_start_reading_input(self.server_sock)
   chunk_type, payload = NailgunProtocol.read_chunk(self.client_sock)
   self.assertEqual(
     (chunk_type, payload),
     (ChunkType.START_READING_INPUT, self.EMPTY_PAYLOAD)
   )

 def test_send_exit(self):
   NailgunProtocol.send_exit(self.server_sock, self.TEST_OUTPUT)
   chunk_type, payload = NailgunProtocol.read_chunk(self.client_sock)
   self.assertEqual(
     (chunk_type, payload),
     (ChunkType.EXIT, self.TEST_OUTPUT)
   )

 def test_send_unicode_chunk(self):
   NailgunProtocol.send_stdout(self.server_sock, self.TEST_UNICODE_PAYLOAD)
   chunk_type, payload = NailgunProtocol.read_chunk(self.client_sock, return_bytes=True)
   self.assertEqual(
     (chunk_type, payload),
     (ChunkType.STDOUT, self.TEST_UNICODE_PAYLOAD)
   )

  def post_fork_child(self):
    """Post-fork child process callback executed via ProcessManager.daemonize()."""
    # Set the Exiter exception hook post-fork so as not to affect the pantsd processes exception
    # hook with socket-specific behavior.
    self._exiter.set_except_hook()

    # Set context in the process title.
    set_process_title('pantsd-runner [{}]'.format(' '.join(self._args)))

    # Broadcast our pid to the remote client so they can send us signals (i.e. SIGINT).
    NailgunProtocol.write_chunk(self._socket, ChunkType.PID, bytes(os.getpid()))

    # Setup a SIGINT signal handler.
    self._setup_sigint_handler()

    # Invoke a Pants run with stdio redirected.
    with self._nailgunned_stdio(self._socket):
      try:
        # Clean global state.
        clean_global_runtime_state(reset_subsystem=True)

        # Re-raise any deferred exceptions, if present.
        self._raise_deferred_exc()

        # Otherwise, conduct a normal run.
        LocalPantsRunner(self._exiter, self._args, self._env, self._graph_helper).run()
      except KeyboardInterrupt:
        self._exiter.exit(1, msg='Interrupted by user.\n')
      except Exception:
        self._exiter.handle_unhandled_exception(add_newline=True)
      else:
        self._exiter.exit(0)

 def test_send_pid(self):
   test_pid = 1
   NailgunProtocol.send_pid(self.server_sock, test_pid)
   chunk_type, payload = NailgunProtocol.read_chunk(self.client_sock, return_bytes=True)
   self.assertEqual(
     (chunk_type, payload),
     (ChunkType.PID, NailgunProtocol.encode_int(test_pid))
   )

  def test_read_chunk_truncated_during_header(self):
    """Construct a chunk and truncate to the first 3 bytes ([:3]), an incomplete header."""
    truncated_chunk = NailgunProtocol.construct_chunk(ChunkType.STDOUT, self.TEST_OUTPUT)[:3]
    self.server_sock.sendall(truncated_chunk)
    self.server_sock.close()

    with self.assertRaises(NailgunProtocol.TruncatedHeaderError):
      NailgunProtocol.read_chunk(self.client_sock)

  def test_read_chunk_truncated_before_payload(self):
    """Construct a chunk and send exactly the header (first 5 bytes) and truncate the remainder."""
    truncated_chunk = NailgunProtocol.construct_chunk(ChunkType.STDOUT, self.TEST_OUTPUT)[:5]
    self.server_sock.sendall(truncated_chunk)
    self.server_sock.close()

    with self.assertRaises(NailgunProtocol.TruncatedPayloadError):
      NailgunProtocol.read_chunk(self.client_sock)

  def test_read_chunk_truncated_during_payload(self):
    """Construct a chunk and truncate the last 3 bytes of the payload ([:-3])."""
    truncated_chunk = NailgunProtocol.construct_chunk(ChunkType.STDOUT, self.TEST_OUTPUT)[:-3]
    self.server_sock.sendall(truncated_chunk)
    self.server_sock.close()

    with self.assertRaises(NailgunProtocol.TruncatedPayloadError):
      NailgunProtocol.read_chunk(self.client_sock)

 def test_send_exit_with_code(self):
   return_code = 1
   NailgunProtocol.send_exit_with_code(self.server_sock, return_code)
   chunk_type, payload = NailgunProtocol.read_chunk(self.client_sock, return_bytes=True)
   self.assertEqual(
     (chunk_type, payload),
     (ChunkType.EXIT, NailgunProtocol.encode_int(return_code))
   )

 def write(self, payload):
   try:
     NailgunProtocol.write_chunk(self._socket, self._chunk_type, payload)
   except IOError as e:
     # If the remote client disconnects and we try to perform a write (e.g. socket.send/sendall),
     # an 'error: [Errno 32] Broken pipe' exception can be thrown. Setting mask_broken_pipe=True
     # safeguards against this case (which is unexpected for most writers of sys.stdout etc) so
     # that we don't awkwardly interrupt the runtime by throwing this exception on writes to
     # stdout/stderr.
     if e.errno == errno.EPIPE and not self._mask_broken_pipe:
       raise

  def post_fork_child(self):
    """Post-fork child process callback executed via ProcessManager.daemonize()."""
    # Set the Exiter exception hook post-fork so as not to affect the pantsd processes exception
    # hook with socket-specific behavior. Note that this intentionally points the faulthandler
    # trace stream to sys.stderr, which at this point is still a _LoggerStream object writing to
    # the `pantsd.log`. This ensures that in the event of e.g. a hung but detached pantsd-runner
    # process that the stacktrace output lands deterministically in a known place vs to a stray
    # terminal window.
    self._exiter.set_except_hook(sys.stderr)

    # Ensure anything referencing sys.argv inherits the Pailgun'd args.
    sys.argv = self._args

    # Set context in the process title.
    set_process_title('pantsd-runner [{}]'.format(' '.join(self._args)))

    # Setup a SIGINT signal handler.
    self._setup_sigint_handler()

    # Broadcast our process group ID (in PID form - i.e. negated) to the remote client so
    # they can send signals (e.g. SIGINT) to all processes in the runners process group.
    pid = str(os.getpgrp() * -1).encode('ascii')
    NailgunProtocol.send_pid(self._socket, pid)

    # Invoke a Pants run with stdio redirected and a proxied environment.
    with self.nailgunned_stdio(self._socket, self._env) as finalizer,\
         hermetic_environment_as(**self._env):
      try:
        # Setup the Exiter's finalizer.
        self._exiter.set_finalizer(finalizer)

        # Clean global state.
        clean_global_runtime_state(reset_subsystem=True)

        # Re-raise any deferred exceptions, if present.
        self._raise_deferred_exc()

        # Otherwise, conduct a normal run.
        runner = LocalPantsRunner.create(
          self._exiter,
          self._args,
          self._env,
          self._target_roots,
          self._graph_helper,
          self._options_bootstrapper
        )
        runner.set_start_time(self._maybe_get_client_start_time_from_env(self._env))
        runner.run()
      except KeyboardInterrupt:
        self._exiter.exit(1, msg='Interrupted by user.\n')
      except Exception:
        self._exiter.handle_unhandled_exception(add_newline=True)
      else:
        self._exiter.exit(0)

  def test_read_and_write_chunk(self):
    # Write a command chunk to the server socket.
    NailgunProtocol.write_chunk(self.server_sock, ChunkType.COMMAND, self.TEST_COMMAND)

    # Read the chunk from the client socket.
    chunk_type, payload = NailgunProtocol.read_chunk(self.client_sock)

    self.assertEqual(
      (chunk_type, payload),
      (ChunkType.COMMAND, self.TEST_COMMAND)
    )