blob: 3a0fbfad90f25b39aad147d046ee4d2a2c0d085b [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.
Simon Glassebec58f2016-07-04 11:58:39 -060021
22 Members:
23 output: accumulated output from expect()
Stephen Warrene8debf32016-01-26 13:41:30 -070024 """
Stephen Warrend2015062016-01-15 11:15:24 -070025
Stephen Warrend27f2fc2016-01-27 23:57:53 -070026 def __init__(self, args, cwd=None):
Stephen Warrene8debf32016-01-26 13:41:30 -070027 """Spawn (fork/exec) the sub-process.
Stephen Warrend2015062016-01-15 11:15:24 -070028
29 Args:
Stephen Warrend27f2fc2016-01-27 23:57:53 -070030 args: array of processs arguments. argv[0] is the command to
31 execute.
32 cwd: the directory to run the process in, or None for no change.
Stephen Warrend2015062016-01-15 11:15:24 -070033
34 Returns:
35 Nothing.
Stephen Warrene8debf32016-01-26 13:41:30 -070036 """
Stephen Warrend2015062016-01-15 11:15:24 -070037
38 self.waited = False
39 self.buf = ''
Simon Glassebec58f2016-07-04 11:58:39 -060040 self.output = ''
Stephen Warrend2015062016-01-15 11:15:24 -070041 self.logfile_read = None
42 self.before = ''
43 self.after = ''
44 self.timeout = None
Stephen Warren085e64d2016-07-06 10:34:30 -060045 # http://stackoverflow.com/questions/7857352/python-regex-to-match-vt100-escape-sequences
46 # Note that re.I doesn't seem to work with this regex (or perhaps the
47 # version of Python in Ubuntu 14.04), hence the inclusion of a-z inside
48 # [] instead.
49 self.re_vt100 = re.compile('(\x1b\[|\x9b)[^@-_a-z]*[@-_a-z]|\x1b[@-_a-z]')
Stephen Warrend2015062016-01-15 11:15:24 -070050
51 (self.pid, self.fd) = pty.fork()
52 if self.pid == 0:
53 try:
54 # For some reason, SIGHUP is set to SIG_IGN at this point when
55 # run under "go" (www.go.cd). Perhaps this happens under any
56 # background (non-interactive) system?
57 signal.signal(signal.SIGHUP, signal.SIG_DFL)
Stephen Warrend27f2fc2016-01-27 23:57:53 -070058 if cwd:
59 os.chdir(cwd)
Stephen Warrend2015062016-01-15 11:15:24 -070060 os.execvp(args[0], args)
61 except:
62 print 'CHILD EXECEPTION:'
63 import traceback
64 traceback.print_exc()
65 finally:
66 os._exit(255)
67
Stephen Warren93134e12016-02-10 16:54:37 -070068 try:
69 self.poll = select.poll()
70 self.poll.register(self.fd, select.POLLIN | select.POLLPRI | select.POLLERR | select.POLLHUP | select.POLLNVAL)
71 except:
72 self.close()
73 raise
Stephen Warrend2015062016-01-15 11:15:24 -070074
75 def kill(self, sig):
Stephen Warrene8debf32016-01-26 13:41:30 -070076 """Send unix signal "sig" to the child process.
Stephen Warrend2015062016-01-15 11:15:24 -070077
78 Args:
79 sig: The signal number to send.
80
81 Returns:
82 Nothing.
Stephen Warrene8debf32016-01-26 13:41:30 -070083 """
Stephen Warrend2015062016-01-15 11:15:24 -070084
85 os.kill(self.pid, sig)
86
87 def isalive(self):
Stephen Warrene8debf32016-01-26 13:41:30 -070088 """Determine whether the child process is still running.
Stephen Warrend2015062016-01-15 11:15:24 -070089
90 Args:
91 None.
92
93 Returns:
94 Boolean indicating whether process is alive.
Stephen Warrene8debf32016-01-26 13:41:30 -070095 """
Stephen Warrend2015062016-01-15 11:15:24 -070096
97 if self.waited:
98 return False
99
100 w = os.waitpid(self.pid, os.WNOHANG)
101 if w[0] == 0:
102 return True
103
104 self.waited = True
105 return False
106
107 def send(self, data):
Stephen Warrene8debf32016-01-26 13:41:30 -0700108 """Send data to the sub-process's stdin.
Stephen Warrend2015062016-01-15 11:15:24 -0700109
110 Args:
111 data: The data to send to the process.
112
113 Returns:
114 Nothing.
Stephen Warrene8debf32016-01-26 13:41:30 -0700115 """
Stephen Warrend2015062016-01-15 11:15:24 -0700116
117 os.write(self.fd, data)
118
119 def expect(self, patterns):
Stephen Warrene8debf32016-01-26 13:41:30 -0700120 """Wait for the sub-process to emit specific data.
Stephen Warrend2015062016-01-15 11:15:24 -0700121
122 This function waits for the process to emit one pattern from the
123 supplied list of patterns, or for a timeout to occur.
124
125 Args:
126 patterns: A list of strings or regex objects that we expect to
127 see in the sub-process' stdout.
128
129 Returns:
130 The index within the patterns array of the pattern the process
131 emitted.
132
133 Notable exceptions:
134 Timeout, if the process did not emit any of the patterns within
135 the expected time.
Stephen Warrene8debf32016-01-26 13:41:30 -0700136 """
Stephen Warrend2015062016-01-15 11:15:24 -0700137
138 for pi in xrange(len(patterns)):
139 if type(patterns[pi]) == type(''):
140 patterns[pi] = re.compile(patterns[pi])
141
Stephen Warrend314e242016-01-22 12:30:07 -0700142 tstart_s = time.time()
Stephen Warrend2015062016-01-15 11:15:24 -0700143 try:
144 while True:
145 earliest_m = None
146 earliest_pi = None
147 for pi in xrange(len(patterns)):
148 pattern = patterns[pi]
149 m = pattern.search(self.buf)
150 if not m:
151 continue
Stephen Warren44ac7622016-01-27 23:57:47 -0700152 if earliest_m and m.start() >= earliest_m.start():
Stephen Warrend2015062016-01-15 11:15:24 -0700153 continue
154 earliest_m = m
155 earliest_pi = pi
156 if earliest_m:
157 pos = earliest_m.start()
Stephen Warrend8926812016-02-05 18:04:42 -0700158 posafter = earliest_m.end()
Stephen Warrend2015062016-01-15 11:15:24 -0700159 self.before = self.buf[:pos]
160 self.after = self.buf[pos:posafter]
Simon Glassebec58f2016-07-04 11:58:39 -0600161 self.output += self.buf[:posafter]
Stephen Warrend2015062016-01-15 11:15:24 -0700162 self.buf = self.buf[posafter:]
163 return earliest_pi
Stephen Warrend314e242016-01-22 12:30:07 -0700164 tnow_s = time.time()
Stephen Warren89ab8412016-02-04 16:11:50 -0700165 if self.timeout:
166 tdelta_ms = (tnow_s - tstart_s) * 1000
167 poll_maxwait = self.timeout - tdelta_ms
168 if tdelta_ms > self.timeout:
169 raise Timeout()
170 else:
171 poll_maxwait = None
172 events = self.poll.poll(poll_maxwait)
Stephen Warrend2015062016-01-15 11:15:24 -0700173 if not events:
174 raise Timeout()
175 c = os.read(self.fd, 1024)
176 if not c:
177 raise EOFError()
178 if self.logfile_read:
179 self.logfile_read.write(c)
180 self.buf += c
Stephen Warren085e64d2016-07-06 10:34:30 -0600181 # count=0 is supposed to be the default, which indicates
182 # unlimited substitutions, but in practice the version of
183 # Python in Ubuntu 14.04 appears to default to count=2!
184 self.buf = self.re_vt100.sub('', self.buf, count=1000000)
Stephen Warrend2015062016-01-15 11:15:24 -0700185 finally:
186 if self.logfile_read:
187 self.logfile_read.flush()
188
189 def close(self):
Stephen Warrene8debf32016-01-26 13:41:30 -0700190 """Close the stdio connection to the sub-process.
Stephen Warrend2015062016-01-15 11:15:24 -0700191
192 This also waits a reasonable time for the sub-process to stop running.
193
194 Args:
195 None.
196
197 Returns:
198 Nothing.
Stephen Warrene8debf32016-01-26 13:41:30 -0700199 """
Stephen Warrend2015062016-01-15 11:15:24 -0700200
201 os.close(self.fd)
202 for i in xrange(100):
203 if not self.isalive():
204 break
205 time.sleep(0.1)
Simon Glassebec58f2016-07-04 11:58:39 -0600206
207 def get_expect_output(self):
208 """Return the output read by expect()
209
210 Returns:
211 The output processed by expect(), as a string.
212 """
213 return self.output