text: # emacs: -*- mode: python; py-indent-offset: 4; tab-width: 4; indent-tabs-mode: nil -*-
  # ex: set sts=4 ts=4 sw=4 et:
  # ## ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ##
  #
  #   See COPYING file distributed along with the datalad package for the
  #   copyright and license terms.
  #
  # ## ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ##
  """
  Thread based subprocess execution with stdout and stderr passed to protocol objects
  """
  
  import enum
  import logging
  import subprocess
  import time
  from collections import deque
  from collections.abc import Generator
  from queue import (
      Empty,
      Queue,
  )
  from typing import (
      Any,
      Dict,
      IO,
      List,
      Optional,
      Type,
      Union,
  )
  
  from datalad.utils import on_windows
  
  from .exception import CommandError
  from .protocol import (
      GeneratorMixIn,
      WitlessProtocol,
  )
  from .runnerthreads import (
      _try_close,
      IOState,
      ReadThread,
      WaitThread,
      WriteThread,
  )
  
  
  lgr = logging.getLogger("datalad.runner.nonasyncrunner")
  
  STDIN_FILENO = 0
  STDOUT_FILENO = 1
  STDERR_FILENO = 2
  
  
  class _ResultGenerator(Generator):
      """
      Generator returned by run_command if the protocol class
      is a subclass of `datalad.runner.protocol.GeneratorMixIn`
      """
      class GeneratorState(enum.Enum):
          process_running = 0
          process_exited = 1
          connection_lost = 2
          waiting_for_process = 3
  
      def __init__(self,
                   runner: "ThreadedRunner",
                   result_queue: deque
                   ):
  
          super().__init__()
          self.runner = runner
          self.result_queue = result_queue
          self.return_code = None
          self.state = self.GeneratorState.process_running
          self.all_closed = False
  
      def _check_result(self):
          self.runner._check_result()
  
      def send(self, _):
          runner = self.runner
          if self.state == self.GeneratorState.process_running:
  
              # If we have elements in the result queue, return one
              while len(self.result_queue) == 0 and runner.should_continue():
                  runner.process_queue()
              if len(self.result_queue) > 0:
                  return self.result_queue.popleft()
  
              # The process must have exited
              # Let the protocol prepare the result. This has to be done after
              # the loop was left to ensure that all data from stdout and stderr
              # is processed.
              runner.protocol.process_exited()
              self.return_code = runner.process.poll()
              self._check_result()
              self.state = self.GeneratorState.process_exited
  
          if self.state == self.GeneratorState.process_exited:
              # The protocol might have added result in the
              # _prepare_result()- or in the process_exited()-
              # callback. Those are returned here.
              if len(self.result_queue) > 0:
                  return self.result_queue.popleft()
  
              runner.ensure_stdin_stdout_stderr_closed()
              runner.protocol.connection_lost(None)   # TODO: check for exceptions
              runner.wait_for_threads()
              self.state = self.GeneratorState.connection_lost
  
          if self.state == self.GeneratorState.connection_lost:
              # Get all results that were enqueued in
              # state: GeneratorState.process_exited.
              if len(self.result_queue) > 0:
                  return self.result_queue.popleft()
              raise StopIteration(self.return_code)
  
      def throw(self, exception_type, value=None, trace_back=None):
          return Generator.throw(self, exception_type, value, trace_back)
  
  
  class ThreadedRunner:
      """
      A class the contains a naive implementation for concurrent sub-process
      execution. It uses `subprocess.Popen` and threads to read from stdout and
      stderr of the subprocess, and to write to stdin of the subprocess.
  
      All read data and timeouts are passed to a protocol instance, which can
      create the final result.
      """
      # Interval in seconds after which we check that a subprocess
      # is still running.
      timeout_resolution = 0.2
  
      def __init__(self,
                   cmd: Union[str, List],
                   protocol_class: Type[WitlessProtocol],
                   stdin: Any,
                   protocol_kwargs: Optional[Dict] = None,
                   timeout: Optional[float] = None,
                   exception_on_error: bool = True,
                   **popen_kwargs
                   ):
          """
          Parameters
          ----------
          cmd : list or str
              Command to be executed, passed to `subprocess.Popen`. If cmd
              is a str, `subprocess.Popen will be called with `shell=True`.
  
          protocol : WitlessProtocol class or subclass which will be
              instantiated for managing communication with the subprocess.
  
              If the protocol is a subclass of
              `datalad.runner.protocol.GeneratorMixIn`, this function will
              return a `Generator` which yields whatever the protocol callback
              fed into `GeneratorMixIn.send_result()`.
  
              If the protocol is not a subclass of
              `datalad.runner.protocol.GeneratorMixIn`, the function will return
              the result created by the protocol method `_generate_result`.
  
          stdin : file-like, string, bytes, Queue, or None
              If stdin is a file-like, it will be directly used as stdin for the
              subprocess. The caller is responsible for writing to it and closing
              it. If stdin is a string or bytes, those will be fed to stdin of the
              subprocess. If all data is written, stdin will be closed.
              If stdin is a Queue, all elements (bytes) put into the Queue will
              be passed to stdin until None is read from the queue. If None is
              read, stdin of the subprocess is closed.
              If stdin is None, nothing will be sent to stdin of the subprocess.
              More precisely, `subprocess.Popen` will be called with `stdin=None`.
  
          protocol_kwargs : dict, optional
              Passed to the protocol class constructor.
  
          timeout : float, optional
              If a non-`None` timeout is specified, the `timeout`-method of
              the protocol will be called if:
  
              - stdin-write, stdout-read, or stderr-read time out. In this case
                the file descriptor will be given as argument to the
                timeout-method. If the timeout-method return `True`, the file
                descriptor will be closed.
  
              - process.wait() timeout: if waiting for process completion after
                stdin, stderr, and stdout takes longer than `timeout` seconds,
                the timeout-method will be called with the argument `None`. If
                it returns `True`, the process will be terminated.
  
          exception_on_error : bool, optional
              This argument is only interpreted if the protocol is a subclass
              of `GeneratorMixIn`. If it is `True` (default), a
              `CommandErrorException` is raised by the generator if the
              sub process exited with a return code not equal to zero. If the
              parameter is `False`, no exception is raised. In both cases the
              return code can be read from the attribute `return_code` of
              the generator.
  
          popen_kwargs : dict, optional
              Passed to `subprocess.Popen`, will typically be parameters
              supported by `subprocess.Popen`. Note that `bufsize`, `stdin`,
              `stdout`, `stderr`, and `shell` will be overwritten internally.
          """
  
          self.cmd = cmd
          self.protocol_class = protocol_class
          self.stdin = stdin
          self.protocol_kwargs = protocol_kwargs or {}
          self.timeout = timeout
          self.exception_on_error = exception_on_error
          self.popen_kwargs = popen_kwargs
  
          self.catch_stdout = self.protocol_class.proc_out is not None
          self.catch_stderr = self.protocol_class.proc_err is not None
          self.generator = self.protocol_class.generator is not None
  
          self.write_stdin = False
          self.stdin_queue = None
          self.protocol = None
          self.process = None
          self.process_stdin_fileno = None
          self.process_stdout_fileno = None
          self.process_stderr_fileno = None
          self.stderr_enqueueing_thread = None
          self.stdout_enqueueing_thread = None
          self.stdin_enqueueing_thread = None
          self.process_waiting_thread = None
  
          self.process_running = False
          self.fileno_mapping = None
          self.fileno_to_file = None
          self.file_to_fileno = None
          self.output_queue = Queue()
          self.result = None
          self.process_removed = False
          self.return_code = None
  
          self.last_touched = dict()
          self.active_file_numbers = set()
          self.stall_check_interval = 10
  
      def _check_result(self):
          if self.exception_on_error is True:
              if self.return_code != 0:
                  protocol = self.protocol
                  decoded_output = {
                      source: protocol.fd_infos[fileno][1].decode(protocol.encoding)
                      for source, fileno in (
                          ("stdout", protocol.stdout_fileno),
                          ("stderr", protocol.stderr_fileno))
                      if protocol.fd_infos[fileno][1] is not None
                  }
                  raise CommandError(cmd=self.cmd,
                                     code=self.return_code,
                                     stdout=decoded_output.get("stdout", None),
                                     stderr=decoded_output.get("stderr", None))
  
      def run(self) -> Union[Any, Generator]:
          """
          Run the command as specified in __init__.
  
          Returns
          -------
          Any
              If the protocol is not a subclass of `GeneratorMixIn`, the
              result of protocol._prepare_result will be returned.
  
          Generator
              If the protocol is a subclass of `GeneratorMixIn`, a Generator
              will be returned. This allows to use this method in constructs
              like:
  
                  for protocol_output in runner.run():
                      ...
  
              Where the iterator yields whatever protocol.pipe_data_received
              sends into the generator.
              If all output was yielded and the process has terminated, the
              generator will raise StopIteration(return_code), where
              return_code is the return code of the process. The return code
              of the process will also be stored in the "return_code"-attribute
              of the runner. So you could write:
  
                 gen = runner.run()
                 for file_descriptor, data in gen:
                     ...
  
                 # get the return code of the process
                 result = gen.return_code
          """
          if isinstance(self.stdin, (int, IO, type(None))):
              # We will not write anything to stdin. If the caller passed a
              # file-like he can write to it from a different thread.
              self.write_stdin = False
  
          elif isinstance(self.stdin, (str, bytes)):
              # Establish a queue to write to the process and
              # enqueue the input that is already provided.
              self.write_stdin = True
              self.stdin_queue = Queue()
              self.stdin_queue.put(self.stdin)
              self.stdin_queue.put(None)
  
          elif isinstance(self.stdin, Queue):
              # Establish a queue to write to the process.
              self.write_stdin = True
              self.stdin_queue = self.stdin
  
          else:
              # We do not recognize the input class will and just pass is through
              # to Popen(). We assume that the caller handles any writing if
              # desired.
              self.write_stdin = False
  
          self.protocol = self.protocol_class(**self.protocol_kwargs)
  
          # The following command is generated internally by datalad
          # and trusted. Security check is therefore skipped.
          kwargs = {
              **self.popen_kwargs,
              **dict(
                  bufsize=0,
                  stdin=subprocess.PIPE if self.write_stdin else self.stdin,
                  stdout=subprocess.PIPE if self.catch_stdout else None,
                  stderr=subprocess.PIPE if self.catch_stderr else None,
                  shell=True if isinstance(self.cmd, str) else False      # nosec
              )
          }
  
          try:
              # The following command is generated internally by datalad
              # and trusted. Security check is therefore skipped.
              self.process = subprocess.Popen(self.cmd, **kwargs)         # nosec
  
          except OSError as e:
              if not on_windows and "argument list too long" in str(e).lower():
                  lgr.error(
                      "Caught exception suggesting too large stack size limits. "
                      "Hint: use 'ulimit -s' command to see current limit and "
                      "e.g. 'ulimit -s 8192' to reduce it to avoid this "
                      "exception. See "
                      "https://github.com/datalad/datalad/issues/6106 for more "
                      "information."
                  )
              raise
          self.process_running = True
          self.active_file_numbers.add(None)
  
          self.process_stdin_fileno = self.process.stdin.fileno() if self.write_stdin else None
          self.process_stdout_fileno = self.process.stdout.fileno() if self.catch_stdout else None
          self.process_stderr_fileno = self.process.stderr.fileno() if self.catch_stderr else None
  
          # We pass process as transport-argument. It does not have the same
          # semantics as the asyncio-signature, but since it is only used in
          # WitlessProtocol, all necessary changes can be made there.
          self.protocol.connection_made(self.process)
  
          # Map the pipe file numbers to stdout and stderr file number, because
          # the latter are hardcoded in the protocol code
          self.fileno_mapping = {
              self.process_stdout_fileno: STDOUT_FILENO,
              self.process_stderr_fileno: STDERR_FILENO,
              self.process_stdin_fileno: STDIN_FILENO,
          }
  
          self.fileno_to_file = {
              self.process_stdout_fileno: self.process.stdout,
              self.process_stderr_fileno: self.process.stderr,
              self.process_stdin_fileno: self.process.stdin
          }
  
          self.file_to_fileno = {
              f: f.fileno()
              for f in (
                  self.process.stdout,
                  self.process.stderr,
                  self.process.stdin
              ) if f is not None
          }
  
          current_time = time.time()
          if self.timeout:
              self.last_touched[None] = current_time
  
          cmd_string = self.cmd if isinstance(self.cmd, str) else " ".join(self.cmd)
          if self.catch_stderr:
              if self.timeout:
                  self.last_touched[self.process_stderr_fileno] = current_time
              self.active_file_numbers.add(self.process_stderr_fileno)
              self.last_touched[self.process_stderr_fileno] = current_time
              self.stderr_enqueueing_thread = ReadThread(
                  identifier="STDERR: " + cmd_string[:20],
                  signal_queues=[self.output_queue],
                  user_info=self.process_stderr_fileno,
                  source=self.process.stderr,
                  destination_queue=self.output_queue)
              self.stderr_enqueueing_thread.start()
  
          if self.catch_stdout:
              if self.timeout:
                  self.last_touched[self.process_stdout_fileno] = current_time
              self.active_file_numbers.add(self.process_stdout_fileno)
              self.last_touched[self.process_stdout_fileno] = current_time
              self.stdout_enqueueing_thread = ReadThread(
                  identifier="STDOUT: " + cmd_string[:20],
                  signal_queues=[self.output_queue],
                  user_info=self.process_stdout_fileno,
                  source=self.process.stdout,
                  destination_queue=self.output_queue)
              self.stdout_enqueueing_thread.start()
  
          if self.write_stdin:
              # No timeouts for stdin
              self.active_file_numbers.add(self.process_stdin_fileno)
              self.stdin_enqueueing_thread = WriteThread(
                  identifier="STDIN: " + cmd_string[:20],
                  user_info=self.process_stdin_fileno,
                  signal_queues=[self.output_queue],
                  source_queue=self.stdin_queue,
                  destination=self.process.stdin)
              self.stdin_enqueueing_thread.start()
  
          self.process_waiting_thread = WaitThread(
              "process_waiter",
              [self.output_queue],
              self.process)
          self.process_waiting_thread.start()
  
          if issubclass(self.protocol_class, GeneratorMixIn):
              return _ResultGenerator(self, self.protocol.result_queue)
  
          return self.process_loop()
  
      def process_loop(self) -> Any:
          # Process internal messages until no more active file descriptors
          # are present. This works because active file numbers are only
          # removed when an EOF is received in `self.process_queue`.
          while self.should_continue():
              self.process_queue()
  
          # Let the protocol prepare the result. This has to be done after
          # the loop was left to ensure that all data from stdout and stderr
          # is processed.
          self.result = self.protocol._prepare_result()
          self.protocol.process_exited()
  
          # Ensure that all communication channels are closed.
          self.ensure_stdin_stdout_stderr_closed()
          self.protocol.connection_lost(None)  # TODO: check exception
          self.wait_for_threads()
          return self.result
  
      def _handle_file_timeout(self, source):
          if self.protocol.timeout(self.fileno_mapping[source]) is True:
              self.remove_file_number(source)
  
      def _handle_process_timeout(self):
          if self.protocol.timeout(None) is True:
              self.ensure_stdin_stdout_stderr_closed()
              self.process.terminate()
              self.process.wait()
              self.remove_process()
  
      def _handle_source_timeout(self, source):
          if source is None:
              self._handle_process_timeout()
          else:
              self._handle_file_timeout(source)
  
      def _update_timeouts(self):
          last_touched = list(self.last_touched.items())
          new_times = dict()
          current_time = time.time()
          for source, last_time in last_touched:
              if current_time - last_time >= self.timeout:
                  new_times[source] = current_time
                  self._handle_source_timeout(source)
          self.last_touched = {
              **self.last_touched,
              **new_times}
  
      def process_timeouts(self):
          if self.timeout is not None:
              self._update_timeouts()
  
      def should_continue(self) -> bool:
          # Continue with queue processing if there is still a process or
          # monitored files, or if there are still elements in the output queue.
          return (
              len(self.active_file_numbers) > 0
              or not self.output_queue.empty())
  
      def is_stalled(self) -> bool:
          # If all threads have exited and the queue is empty, we might have a
          # stall condition.
          live_threads = [
              thread.is_alive()
              for thread in (
                  self.stdout_enqueueing_thread,
                  self.stderr_enqueueing_thread,
                  self.process_waiting_thread,
              ) if thread is not None]
          return not any(live_threads) and self.output_queue.empty()
  
      def check_for_stall(self) -> bool:
          if self.stall_check_interval == 0:
              self.stall_check_interval = 11
              if self.is_stalled():
                  lgr.warning(
                      "ThreadedRunner.process_queue(): stall detected")
                  return True
          self.stall_check_interval -= 1
          return False
  
      def process_queue(self):
          """
          Get a single event from the queue or handle a timeout. This method
          might modify the set of active file numbers if a file-closed event
          is read from the output queue, or if a timeout-callback return True.
          """
          data = None
          while True:
              # We do not need a user provided timeout here. If
              # self.timeout is None, no timeouts are reported anyway.
              # If self.timeout is not None, and any enqueuing (stdin)
              # or de-queuing (stdout, stderr) operation takes longer than
              # self.timeout, we will get a queue entry for that.
              # We still use a "system"-timeout, i.e.
              # `ThreadedRunner.process_check_interval`, to check whether the
              # process is still running.
              try:
                  file_number, state, data = self.output_queue.get(
                      timeout=ThreadedRunner.timeout_resolution)
                  break
              except Empty:
                  if self.check_for_stall() is True:
                      return
                  self.process_timeouts()
                  continue
  
          if state == IOState.process_exit:
              self.remove_process()
              return
  
          if self.write_stdin and file_number == self.process_stdin_fileno:
              # The only data-signal we expect from stdin thread
              # is None, indicating that the thread ended
              assert data is None
              self.remove_file_number(self.process_stdin_fileno)
  
          elif self.catch_stderr or self.catch_stdout:
              if data is None:
                  # Received an EOF for stdout or stderr.
                  self.remove_file_number(file_number)
              else:
                  # Call the protocol handler for data
                  assert isinstance(data, bytes)
                  self.last_touched[file_number] = time.time()
                  self.protocol.pipe_data_received(
                      self.fileno_mapping[file_number],
                      data)
  
      def remove_process(self):
          if None not in self.active_file_numbers:
              # Might already be removed due to a timeout callback returning
              # True and subsequent removal of the process.
              return
          self.active_file_numbers.remove(None)
          if self.timeout:
              del self.last_touched[None]
  
          # Remove stdin from the active set because the process will
          # no longer consume input from stdin. This is done by enqueuing
          # None to the stdin queue.
          if self.write_stdin:
              self.stdin_queue.put(None)
  
          self.return_code = self.process.poll()
  
      def remove_file_number(self, file_number: int):
          """
          Remove a file number from the active set and from
          the timeout set.
          """
  
          # TODO: check exception
          # Let the protocol know that the connection was lost.
          self.protocol.pipe_connection_lost(
              self.fileno_mapping[file_number],
              None)
  
          if file_number in self.active_file_numbers:
              # Remove the file number from the set of active numbers.
              self.active_file_numbers.remove(file_number)
  
          # If we are checking timeouts, remove the file number from
          # timeouts.
          if self.timeout and file_number in self.last_touched:
              del self.last_touched[file_number]
  
          _try_close(self.fileno_to_file[file_number])
  
      def close_stdin(self):
          if self.stdin_queue:
              self.stdin_queue.put(None)
  
      def _ensure_closed(self, file_objects):
          for file_object in file_objects:
              if file_object is not None:
                  file_number = self.file_to_fileno.get(file_object, None)
                  if file_number is not None:
                      if self.timeout and file_number in self.last_touched:
                          del self.last_touched[file_number]
                      if file_number in self.active_file_numbers:
                          self.active_file_numbers.remove(file_number)
                  _try_close(file_object)
  
      def ensure_stdin_stdout_stderr_closed(self):
          self.close_stdin()
          self._ensure_closed((self.process.stdin,
                               self.process.stdout,
                               self.process.stderr))
  
      def ensure_stdout_stderr_closed(self):
          self._ensure_closed((self.process.stdout, self.process.stderr))
  
      def wait_for_threads(self):
          for thread in (self.stderr_enqueueing_thread,
                         self.stdout_enqueueing_thread,
                         self.stdin_enqueueing_thread):
              if thread is not None:
                  thread.request_exit()
  
  
  def run_command(cmd: Union[str, List],
                  protocol: Type[WitlessProtocol],
                  stdin: Any,
                  protocol_kwargs: Optional[Dict] = None,
                  timeout: Optional[float] = None,
                  exception_on_error: bool = True,
                  **popen_kwargs) -> Union[Any, Generator]:
      """
      Run a command in a subprocess
  
      this function delegates the execution to an instance of
      `ThreadedRunner`, please see `ThreadedRunner.__init__()` for a
      documentation of the parameters, and `ThreadedRunner.run()` for a
      documentation of the return values.
      """
      runner = ThreadedRunner(
          cmd=cmd,
          protocol_class=protocol,
          stdin=stdin,
          protocol_kwargs=protocol_kwargs,
          timeout=timeout,
          exception_on_error=exception_on_error,
          **popen_kwargs,
      )
  
      return runner.run()
  