blob: 4b9e81af3efbdd8e2d96db58273f194906ef91e6 [file] [log] [blame]
Stephen Warrend2015062016-01-15 11:15:24 -07001# Copyright (c) 2015-2016, NVIDIA CORPORATION. All rights reserved.
2#
3# SPDX-License-Identifier: GPL-2.0
4
5# Logic to spawn a sub-process and interact with its stdio.
6
7import os
8import re
9import pty
10import signal
11import select
12import time
13
14class Timeout(Exception):
Stephen Warrene8debf32016-01-26 13:41:30 -070015 """An exception sub-class that indicates that a timeout occurred."""
Stephen Warrend2015062016-01-15 11:15:24 -070016 pass
17
18class Spawn(object):
Stephen Warrene8debf32016-01-26 13:41:30 -070019 """Represents the stdio of a freshly created sub-process. Commands may be
Stephen Warrend2015062016-01-15 11:15:24 -070020 sent to the process, and responses waited for.
Stephen Warrene8debf32016-01-26 13:41:30 -070021 """
Stephen Warrend2015062016-01-15 11:15:24 -070022
Stephen Warrend27f2fc2016-01-27 23:57:53 -070023 def __init__(self, args, cwd=None):
Stephen Warrene8debf32016-01-26 13:41:30 -070024 """Spawn (fork/exec) the sub-process.
Stephen Warrend2015062016-01-15 11:15:24 -070025
26 Args:
Stephen Warrend27f2fc2016-01-27 23:57:53 -070027 args: array of processs arguments. argv[0] is the command to
28 execute.
29 cwd: the directory to run the process in, or None for no change.
Stephen Warrend2015062016-01-15 11:15:24 -070030
31 Returns:
32 Nothing.
Stephen Warrene8debf32016-01-26 13:41:30 -070033 """
Stephen Warrend2015062016-01-15 11:15:24 -070034
35 self.waited = False
36 self.buf = ''
37 self.logfile_read = None
38 self.before = ''
39 self.after = ''
40 self.timeout = None
41
42 (self.pid, self.fd) = pty.fork()
43 if self.pid == 0:
44 try:
45 # For some reason, SIGHUP is set to SIG_IGN at this point when
46 # run under "go" (www.go.cd). Perhaps this happens under any
47 # background (non-interactive) system?
48 signal.signal(signal.SIGHUP, signal.SIG_DFL)
Stephen Warrend27f2fc2016-01-27 23:57:53 -070049 if cwd:
50 os.chdir(cwd)
Stephen Warrend2015062016-01-15 11:15:24 -070051 os.execvp(args[0], args)
52 except:
53 print 'CHILD EXECEPTION:'
54 import traceback
55 traceback.print_exc()
56 finally:
57 os._exit(255)
58
59 self.poll = select.poll()
60 self.poll.register(self.fd, select.POLLIN | select.POLLPRI | select.POLLERR | select.POLLHUP | select.POLLNVAL)
61
62 def kill(self, sig):
Stephen Warrene8debf32016-01-26 13:41:30 -070063 """Send unix signal "sig" to the child process.
Stephen Warrend2015062016-01-15 11:15:24 -070064
65 Args:
66 sig: The signal number to send.
67
68 Returns:
69 Nothing.
Stephen Warrene8debf32016-01-26 13:41:30 -070070 """
Stephen Warrend2015062016-01-15 11:15:24 -070071
72 os.kill(self.pid, sig)
73
74 def isalive(self):
Stephen Warrene8debf32016-01-26 13:41:30 -070075 """Determine whether the child process is still running.
Stephen Warrend2015062016-01-15 11:15:24 -070076
77 Args:
78 None.
79
80 Returns:
81 Boolean indicating whether process is alive.
Stephen Warrene8debf32016-01-26 13:41:30 -070082 """
Stephen Warrend2015062016-01-15 11:15:24 -070083
84 if self.waited:
85 return False
86
87 w = os.waitpid(self.pid, os.WNOHANG)
88 if w[0] == 0:
89 return True
90
91 self.waited = True
92 return False
93
94 def send(self, data):
Stephen Warrene8debf32016-01-26 13:41:30 -070095 """Send data to the sub-process's stdin.
Stephen Warrend2015062016-01-15 11:15:24 -070096
97 Args:
98 data: The data to send to the process.
99
100 Returns:
101 Nothing.
Stephen Warrene8debf32016-01-26 13:41:30 -0700102 """
Stephen Warrend2015062016-01-15 11:15:24 -0700103
104 os.write(self.fd, data)
105
106 def expect(self, patterns):
Stephen Warrene8debf32016-01-26 13:41:30 -0700107 """Wait for the sub-process to emit specific data.
Stephen Warrend2015062016-01-15 11:15:24 -0700108
109 This function waits for the process to emit one pattern from the
110 supplied list of patterns, or for a timeout to occur.
111
112 Args:
113 patterns: A list of strings or regex objects that we expect to
114 see in the sub-process' stdout.
115
116 Returns:
117 The index within the patterns array of the pattern the process
118 emitted.
119
120 Notable exceptions:
121 Timeout, if the process did not emit any of the patterns within
122 the expected time.
Stephen Warrene8debf32016-01-26 13:41:30 -0700123 """
Stephen Warrend2015062016-01-15 11:15:24 -0700124
125 for pi in xrange(len(patterns)):
126 if type(patterns[pi]) == type(''):
127 patterns[pi] = re.compile(patterns[pi])
128
Stephen Warrend314e242016-01-22 12:30:07 -0700129 tstart_s = time.time()
Stephen Warrend2015062016-01-15 11:15:24 -0700130 try:
131 while True:
132 earliest_m = None
133 earliest_pi = None
134 for pi in xrange(len(patterns)):
135 pattern = patterns[pi]
136 m = pattern.search(self.buf)
137 if not m:
138 continue
Stephen Warren44ac7622016-01-27 23:57:47 -0700139 if earliest_m and m.start() >= earliest_m.start():
Stephen Warrend2015062016-01-15 11:15:24 -0700140 continue
141 earliest_m = m
142 earliest_pi = pi
143 if earliest_m:
144 pos = earliest_m.start()
145 posafter = earliest_m.end() + 1
146 self.before = self.buf[:pos]
147 self.after = self.buf[pos:posafter]
148 self.buf = self.buf[posafter:]
149 return earliest_pi
Stephen Warrend314e242016-01-22 12:30:07 -0700150 tnow_s = time.time()
Stephen Warren89ab8412016-02-04 16:11:50 -0700151 if self.timeout:
152 tdelta_ms = (tnow_s - tstart_s) * 1000
153 poll_maxwait = self.timeout - tdelta_ms
154 if tdelta_ms > self.timeout:
155 raise Timeout()
156 else:
157 poll_maxwait = None
158 events = self.poll.poll(poll_maxwait)
Stephen Warrend2015062016-01-15 11:15:24 -0700159 if not events:
160 raise Timeout()
161 c = os.read(self.fd, 1024)
162 if not c:
163 raise EOFError()
164 if self.logfile_read:
165 self.logfile_read.write(c)
166 self.buf += c
167 finally:
168 if self.logfile_read:
169 self.logfile_read.flush()
170
171 def close(self):
Stephen Warrene8debf32016-01-26 13:41:30 -0700172 """Close the stdio connection to the sub-process.
Stephen Warrend2015062016-01-15 11:15:24 -0700173
174 This also waits a reasonable time for the sub-process to stop running.
175
176 Args:
177 None.
178
179 Returns:
180 Nothing.
Stephen Warrene8debf32016-01-26 13:41:30 -0700181 """
Stephen Warrend2015062016-01-15 11:15:24 -0700182
183 os.close(self.fd)
184 for i in xrange(100):
185 if not self.isalive():
186 break
187 time.sleep(0.1)