| 199 | |
| 200 | |
| 201 | class ClientDriver: |
| 202 | CLIENT_SERVER = os.path.join( |
| 203 | os.path.dirname(os.path.abspath(__file__)), 'cmd-runner' |
| 204 | ) |
| 205 | |
| 206 | def __init__(self): |
| 207 | self._popen = None |
| 208 | self.memory_samples = [] |
| 209 | |
| 210 | def _get_memory_with_ps(self, pid): |
| 211 | # It would be better to eventually switch to psutil, |
| 212 | # which should allow us to test on windows, but for now |
| 213 | # we'll just use ps and run on POSIX platforms. |
| 214 | command_list = ['ps', '-p', str(pid), '-o', 'rss'] |
| 215 | p = Popen(command_list, stdout=PIPE) |
| 216 | stdout = p.communicate()[0] |
| 217 | if not p.returncode == 0: |
| 218 | raise RuntimeError("Could not retrieve memory") |
| 219 | else: |
| 220 | # Get the RSS from output that looks like this: |
| 221 | # RSS |
| 222 | # 4496 |
| 223 | return int(stdout.splitlines()[1].split()[0]) * 1024 |
| 224 | |
| 225 | def record_memory(self): |
| 226 | mem = self._get_memory_with_ps(self._popen.pid) |
| 227 | self.memory_samples.append(mem) |
| 228 | |
| 229 | def start(self, env=None): |
| 230 | """Start up the command runner process.""" |
| 231 | self._popen = Popen( |
| 232 | [sys.executable, self.CLIENT_SERVER], |
| 233 | stdout=PIPE, |
| 234 | stdin=PIPE, |
| 235 | env=env, |
| 236 | ) |
| 237 | |
| 238 | def stop(self): |
| 239 | """Shutdown the command runner process.""" |
| 240 | self.cmd('exit') |
| 241 | self._popen.wait() |
| 242 | |
| 243 | def send_cmd(self, *cmd): |
| 244 | """Send a command and return immediately. |
| 245 | |
| 246 | This is a lower level method than cmd(). |
| 247 | This method will instruct the cmd-runner process |
| 248 | to execute a command, but this method will |
| 249 | immediately return. You will need to use |
| 250 | ``is_cmd_finished()`` to check that the command |
| 251 | is finished. |
| 252 | |
| 253 | This method is useful if you want to record attributes |
| 254 | about the process while an operation is occurring. For |
| 255 | example, if you want to instruct the cmd-runner process |
| 256 | to upload a 1GB file to S3 and you'd like to record |
| 257 | the memory during the upload process, you can use |
| 258 | send_cmd() instead of cmd(). |