]>
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 SW |
22 | |
23 | def __init__(self, args): | |
e8debf39 | 24 | """Spawn (fork/exec) the sub-process. |
d201506c SW |
25 | |
26 | Args: | |
27 | args: array of processs arguments. argv[0] is the command to execute. | |
28 | ||
29 | Returns: | |
30 | Nothing. | |
e8debf39 | 31 | """ |
d201506c SW |
32 | |
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): | |
e8debf39 | 59 | """Send unix signal "sig" to the child process. |
d201506c SW |
60 | |
61 | Args: | |
62 | sig: The signal number to send. | |
63 | ||
64 | Returns: | |
65 | Nothing. | |
e8debf39 | 66 | """ |
d201506c SW |
67 | |
68 | os.kill(self.pid, sig) | |
69 | ||
70 | def isalive(self): | |
e8debf39 | 71 | """Determine whether the child process is still running. |
d201506c SW |
72 | |
73 | Args: | |
74 | None. | |
75 | ||
76 | Returns: | |
77 | Boolean indicating whether process is alive. | |
e8debf39 | 78 | """ |
d201506c SW |
79 | |
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): | |
e8debf39 | 91 | """Send data to the sub-process's stdin. |
d201506c SW |
92 | |
93 | Args: | |
94 | data: The data to send to the process. | |
95 | ||
96 | Returns: | |
97 | Nothing. | |
e8debf39 | 98 | """ |
d201506c SW |
99 | |
100 | os.write(self.fd, data) | |
101 | ||
102 | def expect(self, patterns): | |
e8debf39 | 103 | """Wait for the sub-process to emit specific data. |
d201506c SW |
104 | |
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. | |
e8debf39 | 119 | """ |
d201506c SW |
120 | |
121 | for pi in xrange(len(patterns)): | |
122 | if type(patterns[pi]) == type(''): | |
123 | patterns[pi] = re.compile(patterns[pi]) | |
124 | ||
d314e247 | 125 | tstart_s = time.time() |
d201506c SW |
126 | 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 | |
44ac762b | 135 | if earliest_m and m.start() >= earliest_m.start(): |
d201506c SW |
136 | 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 | |
d314e247 SW |
146 | 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) | |
d201506c SW |
151 | 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): | |
e8debf39 | 164 | """Close the stdio connection to the sub-process. |
d201506c SW |
165 | |
166 | This also waits a reasonable time for the sub-process to stop running. | |
167 | ||
168 | Args: | |
169 | None. | |
170 | ||
171 | Returns: | |
172 | Nothing. | |
e8debf39 | 173 | """ |
d201506c SW |
174 | |
175 | os.close(self.fd) | |
176 | for i in xrange(100): | |
177 | if not self.isalive(): | |
178 | break | |
179 | time.sleep(0.1) |