]>
Commit | Line | Data |
---|---|---|
d201506c SW |
1 | # 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 | ||
7 | import os | |
8 | import re | |
9 | import pty | |
10 | import signal | |
11 | import select | |
12 | import time | |
13 | ||
14 | class Timeout(Exception): | |
e8debf39 | 15 | """An exception sub-class that indicates that a timeout occurred.""" |
d201506c SW |
16 | pass |
17 | ||
18 | class Spawn(object): | |
e8debf39 | 19 | """Represents the stdio of a freshly created sub-process. Commands may be |
d201506c | 20 | sent to the process, and responses waited for. |
e8debf39 | 21 | """ |
d201506c | 22 | |
d27f2fc1 | 23 | def __init__(self, args, cwd=None): |
e8debf39 | 24 | """Spawn (fork/exec) the sub-process. |
d201506c SW |
25 | |
26 | Args: | |
d27f2fc1 SW |
27 | 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. | |
d201506c SW |
30 | |
31 | Returns: | |
32 | Nothing. | |
e8debf39 | 33 | """ |
d201506c SW |
34 | |
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) | |
d27f2fc1 SW |
49 | if cwd: |
50 | os.chdir(cwd) | |
d201506c SW |
51 | 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): | |
e8debf39 | 63 | """Send unix signal "sig" to the child process. |
d201506c SW |
64 | |
65 | Args: | |
66 | sig: The signal number to send. | |
67 | ||
68 | Returns: | |
69 | Nothing. | |
e8debf39 | 70 | """ |
d201506c SW |
71 | |
72 | os.kill(self.pid, sig) | |
73 | ||
74 | def isalive(self): | |
e8debf39 | 75 | """Determine whether the child process is still running. |
d201506c SW |
76 | |
77 | Args: | |
78 | None. | |
79 | ||
80 | Returns: | |
81 | Boolean indicating whether process is alive. | |
e8debf39 | 82 | """ |
d201506c SW |
83 | |
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): | |
e8debf39 | 95 | """Send data to the sub-process's stdin. |
d201506c SW |
96 | |
97 | Args: | |
98 | data: The data to send to the process. | |
99 | ||
100 | Returns: | |
101 | Nothing. | |
e8debf39 | 102 | """ |
d201506c SW |
103 | |
104 | os.write(self.fd, data) | |
105 | ||
106 | def expect(self, patterns): | |
e8debf39 | 107 | """Wait for the sub-process to emit specific data. |
d201506c SW |
108 | |
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. | |
e8debf39 | 123 | """ |
d201506c SW |
124 | |
125 | for pi in xrange(len(patterns)): | |
126 | if type(patterns[pi]) == type(''): | |
127 | patterns[pi] = re.compile(patterns[pi]) | |
128 | ||
d314e247 | 129 | tstart_s = time.time() |
d201506c SW |
130 | 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 | |
44ac762b | 139 | if earliest_m and m.start() >= earliest_m.start(): |
d201506c SW |
140 | 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 | |
d314e247 | 150 | tnow_s = time.time() |
89ab8410 SW |
151 | 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) | |
d201506c SW |
159 | 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): | |
e8debf39 | 172 | """Close the stdio connection to the sub-process. |
d201506c SW |
173 | |
174 | This also waits a reasonable time for the sub-process to stop running. | |
175 | ||
176 | Args: | |
177 | None. | |
178 | ||
179 | Returns: | |
180 | Nothing. | |
e8debf39 | 181 | """ |
d201506c SW |
182 | |
183 | os.close(self.fd) | |
184 | for i in xrange(100): | |
185 | if not self.isalive(): | |
186 | break | |
187 | time.sleep(0.1) |