blob: 7451455671b9b52ab6797949f4de2cfaf712f1a9 [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
23 def __init__(self, args):
Stephen Warrene8debf32016-01-26 13:41:30 -070024 """Spawn (fork/exec) the sub-process.
Stephen Warrend2015062016-01-15 11:15:24 -070025
26 Args:
27 args: array of processs arguments. argv[0] is the command to execute.
28
29 Returns:
30 Nothing.
Stephen Warrene8debf32016-01-26 13:41:30 -070031 """
Stephen Warrend2015062016-01-15 11:15:24 -070032
33 self.waited = False
34 self.buf = ''
35 self.logfile_read = None
36 self.before = ''
37 self.after = ''
38 self.timeout = None
39
40 (self.pid, self.fd) = pty.fork()
41 if self.pid == 0:
42 try:
43 # For some reason, SIGHUP is set to SIG_IGN at this point when
44 # run under "go" (www.go.cd). Perhaps this happens under any
45 # background (non-interactive) system?
46 signal.signal(signal.SIGHUP, signal.SIG_DFL)
47 os.execvp(args[0], args)
48 except:
49 print 'CHILD EXECEPTION:'
50 import traceback
51 traceback.print_exc()
52 finally:
53 os._exit(255)
54
55 self.poll = select.poll()
56 self.poll.register(self.fd, select.POLLIN | select.POLLPRI | select.POLLERR | select.POLLHUP | select.POLLNVAL)
57
58 def kill(self, sig):
Stephen Warrene8debf32016-01-26 13:41:30 -070059 """Send unix signal "sig" to the child process.
Stephen Warrend2015062016-01-15 11:15:24 -070060
61 Args:
62 sig: The signal number to send.
63
64 Returns:
65 Nothing.
Stephen Warrene8debf32016-01-26 13:41:30 -070066 """
Stephen Warrend2015062016-01-15 11:15:24 -070067
68 os.kill(self.pid, sig)
69
70 def isalive(self):
Stephen Warrene8debf32016-01-26 13:41:30 -070071 """Determine whether the child process is still running.
Stephen Warrend2015062016-01-15 11:15:24 -070072
73 Args:
74 None.
75
76 Returns:
77 Boolean indicating whether process is alive.
Stephen Warrene8debf32016-01-26 13:41:30 -070078 """
Stephen Warrend2015062016-01-15 11:15:24 -070079
80 if self.waited:
81 return False
82
83 w = os.waitpid(self.pid, os.WNOHANG)
84 if w[0] == 0:
85 return True
86
87 self.waited = True
88 return False
89
90 def send(self, data):
Stephen Warrene8debf32016-01-26 13:41:30 -070091 """Send data to the sub-process's stdin.
Stephen Warrend2015062016-01-15 11:15:24 -070092
93 Args:
94 data: The data to send to the process.
95
96 Returns:
97 Nothing.
Stephen Warrene8debf32016-01-26 13:41:30 -070098 """
Stephen Warrend2015062016-01-15 11:15:24 -070099
100 os.write(self.fd, data)
101
102 def expect(self, patterns):
Stephen Warrene8debf32016-01-26 13:41:30 -0700103 """Wait for the sub-process to emit specific data.
Stephen Warrend2015062016-01-15 11:15:24 -0700104
105 This function waits for the process to emit one pattern from the
106 supplied list of patterns, or for a timeout to occur.
107
108 Args:
109 patterns: A list of strings or regex objects that we expect to
110 see in the sub-process' stdout.
111
112 Returns:
113 The index within the patterns array of the pattern the process
114 emitted.
115
116 Notable exceptions:
117 Timeout, if the process did not emit any of the patterns within
118 the expected time.
Stephen Warrene8debf32016-01-26 13:41:30 -0700119 """
Stephen Warrend2015062016-01-15 11:15:24 -0700120
121 for pi in xrange(len(patterns)):
122 if type(patterns[pi]) == type(''):
123 patterns[pi] = re.compile(patterns[pi])
124
Stephen Warrend314e242016-01-22 12:30:07 -0700125 tstart_s = time.time()
Stephen Warrend2015062016-01-15 11:15:24 -0700126 try:
127 while True:
128 earliest_m = None
129 earliest_pi = None
130 for pi in xrange(len(patterns)):
131 pattern = patterns[pi]
132 m = pattern.search(self.buf)
133 if not m:
134 continue
Stephen Warren44ac7622016-01-27 23:57:47 -0700135 if earliest_m and m.start() >= earliest_m.start():
Stephen Warrend2015062016-01-15 11:15:24 -0700136 continue
137 earliest_m = m
138 earliest_pi = pi
139 if earliest_m:
140 pos = earliest_m.start()
141 posafter = earliest_m.end() + 1
142 self.before = self.buf[:pos]
143 self.after = self.buf[pos:posafter]
144 self.buf = self.buf[posafter:]
145 return earliest_pi
Stephen Warrend314e242016-01-22 12:30:07 -0700146 tnow_s = time.time()
147 tdelta_ms = (tnow_s - tstart_s) * 1000
148 if tdelta_ms > self.timeout:
149 raise Timeout()
150 events = self.poll.poll(self.timeout - tdelta_ms)
Stephen Warrend2015062016-01-15 11:15:24 -0700151 if not events:
152 raise Timeout()
153 c = os.read(self.fd, 1024)
154 if not c:
155 raise EOFError()
156 if self.logfile_read:
157 self.logfile_read.write(c)
158 self.buf += c
159 finally:
160 if self.logfile_read:
161 self.logfile_read.flush()
162
163 def close(self):
Stephen Warrene8debf32016-01-26 13:41:30 -0700164 """Close the stdio connection to the sub-process.
Stephen Warrend2015062016-01-15 11:15:24 -0700165
166 This also waits a reasonable time for the sub-process to stop running.
167
168 Args:
169 None.
170
171 Returns:
172 Nothing.
Stephen Warrene8debf32016-01-26 13:41:30 -0700173 """
Stephen Warrend2015062016-01-15 11:15:24 -0700174
175 os.close(self.fd)
176 for i in xrange(100):
177 if not self.isalive():
178 break
179 time.sleep(0.1)