u_boot_spawn.py 5.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191
  1. # Copyright (c) 2015-2016, NVIDIA CORPORATION. All rights reserved.
  2. #
  3. # SPDX-License-Identifier: GPL-2.0
  4. # Logic to spawn a sub-process and interact with its stdio.
  5. import os
  6. import re
  7. import pty
  8. import signal
  9. import select
  10. import time
  11. class Timeout(Exception):
  12. """An exception sub-class that indicates that a timeout occurred."""
  13. pass
  14. class Spawn(object):
  15. """Represents the stdio of a freshly created sub-process. Commands may be
  16. sent to the process, and responses waited for.
  17. """
  18. def __init__(self, args, cwd=None):
  19. """Spawn (fork/exec) the sub-process.
  20. Args:
  21. args: array of processs arguments. argv[0] is the command to
  22. execute.
  23. cwd: the directory to run the process in, or None for no change.
  24. Returns:
  25. Nothing.
  26. """
  27. self.waited = False
  28. self.buf = ''
  29. self.logfile_read = None
  30. self.before = ''
  31. self.after = ''
  32. self.timeout = None
  33. (self.pid, self.fd) = pty.fork()
  34. if self.pid == 0:
  35. try:
  36. # For some reason, SIGHUP is set to SIG_IGN at this point when
  37. # run under "go" (www.go.cd). Perhaps this happens under any
  38. # background (non-interactive) system?
  39. signal.signal(signal.SIGHUP, signal.SIG_DFL)
  40. if cwd:
  41. os.chdir(cwd)
  42. os.execvp(args[0], args)
  43. except:
  44. print 'CHILD EXECEPTION:'
  45. import traceback
  46. traceback.print_exc()
  47. finally:
  48. os._exit(255)
  49. try:
  50. self.poll = select.poll()
  51. self.poll.register(self.fd, select.POLLIN | select.POLLPRI | select.POLLERR | select.POLLHUP | select.POLLNVAL)
  52. except:
  53. self.close()
  54. raise
  55. def kill(self, sig):
  56. """Send unix signal "sig" to the child process.
  57. Args:
  58. sig: The signal number to send.
  59. Returns:
  60. Nothing.
  61. """
  62. os.kill(self.pid, sig)
  63. def isalive(self):
  64. """Determine whether the child process is still running.
  65. Args:
  66. None.
  67. Returns:
  68. Boolean indicating whether process is alive.
  69. """
  70. if self.waited:
  71. return False
  72. w = os.waitpid(self.pid, os.WNOHANG)
  73. if w[0] == 0:
  74. return True
  75. self.waited = True
  76. return False
  77. def send(self, data):
  78. """Send data to the sub-process's stdin.
  79. Args:
  80. data: The data to send to the process.
  81. Returns:
  82. Nothing.
  83. """
  84. os.write(self.fd, data)
  85. def expect(self, patterns):
  86. """Wait for the sub-process to emit specific data.
  87. This function waits for the process to emit one pattern from the
  88. supplied list of patterns, or for a timeout to occur.
  89. Args:
  90. patterns: A list of strings or regex objects that we expect to
  91. see in the sub-process' stdout.
  92. Returns:
  93. The index within the patterns array of the pattern the process
  94. emitted.
  95. Notable exceptions:
  96. Timeout, if the process did not emit any of the patterns within
  97. the expected time.
  98. """
  99. for pi in xrange(len(patterns)):
  100. if type(patterns[pi]) == type(''):
  101. patterns[pi] = re.compile(patterns[pi])
  102. tstart_s = time.time()
  103. try:
  104. while True:
  105. earliest_m = None
  106. earliest_pi = None
  107. for pi in xrange(len(patterns)):
  108. pattern = patterns[pi]
  109. m = pattern.search(self.buf)
  110. if not m:
  111. continue
  112. if earliest_m and m.start() >= earliest_m.start():
  113. continue
  114. earliest_m = m
  115. earliest_pi = pi
  116. if earliest_m:
  117. pos = earliest_m.start()
  118. posafter = earliest_m.end()
  119. self.before = self.buf[:pos]
  120. self.after = self.buf[pos:posafter]
  121. self.buf = self.buf[posafter:]
  122. return earliest_pi
  123. tnow_s = time.time()
  124. if self.timeout:
  125. tdelta_ms = (tnow_s - tstart_s) * 1000
  126. poll_maxwait = self.timeout - tdelta_ms
  127. if tdelta_ms > self.timeout:
  128. raise Timeout()
  129. else:
  130. poll_maxwait = None
  131. events = self.poll.poll(poll_maxwait)
  132. if not events:
  133. raise Timeout()
  134. c = os.read(self.fd, 1024)
  135. if not c:
  136. raise EOFError()
  137. if self.logfile_read:
  138. self.logfile_read.write(c)
  139. self.buf += c
  140. finally:
  141. if self.logfile_read:
  142. self.logfile_read.flush()
  143. def close(self):
  144. """Close the stdio connection to the sub-process.
  145. This also waits a reasonable time for the sub-process to stop running.
  146. Args:
  147. None.
  148. Returns:
  149. Nothing.
  150. """
  151. os.close(self.fd)
  152. for i in xrange(100):
  153. if not self.isalive():
  154. break
  155. time.sleep(0.1)