testing/mozharness/mozprocess/pid.py

   1 #!/usr/bin/env python
   2
   3 # This Source Code Form is subject to the terms of the Mozilla Public
   4 # License, v. 2.0. If a copy of the MPL was not distributed with this file,
   5 # You can obtain one at http://mozilla.org/MPL/2.0/.
   6
   7 import os
   8 import mozinfo
   9 import shlex
  10 import subprocess
  11 import sys
  12
  13 # determine the platform-specific invocation of `ps`
  14 if mozinfo.isMac:
  15     psarg = '-Acj'
  16 elif mozinfo.isLinux:
  17     psarg = 'axwww'
  18 else:
  19     psarg = 'ax'
  20
  21 def ps(arg=psarg):
  22     """
  23     python front-end to `ps`
  24     http://en.wikipedia.org/wiki/Ps_%28Unix%29
  25     returns a list of process dicts based on the `ps` header
  26     """
  27     retval = []
  28     process = subprocess.Popen(['ps', arg], stdout=subprocess.PIPE)
  29     stdout, _ = process.communicate()
  30     header = None
  31     for line in stdout.splitlines():
  32         line = line.strip()
  33         if header is None:
  34             # first line is the header
  35             header = line.split()
  36             continue
  37         split = line.split(None, len(header)-1)
  38         process_dict = dict(zip(header, split))
  39         retval.append(process_dict)
  40     return retval
  41
  42 def running_processes(name, psarg=psarg, defunct=True):
  43     """
  44     returns a list of
  45     {'PID': PID of process (int)
  46      'command': command line of process (list)}
  47      with the executable named `name`.
  48      - defunct: whether to return defunct processes
  49     """
  50     retval = []
  51     for process in ps(psarg):
  52         # Support for both BSD and UNIX syntax
  53         # `ps aux` returns COMMAND, `ps -ef` returns CMD
  54         try:
  55             command = process['COMMAND']
  56         except KeyError:
  57             command = process['CMD']
  58
  59         command = shlex.split(command)
  60         if command[-1] == '<defunct>':
  61             command = command[:-1]
  62             if not command or not defunct:
  63                 continue
  64         if 'STAT' in process and not defunct:
  65             if process['STAT'] == 'Z+':
  66                 continue
  67         prog = command[0]
  68         basename = os.path.basename(prog)
  69         if basename == name:
  70             retval.append((int(process['PID']), command))
  71     return retval
  72
  73 def get_pids(name):
  74     """Get all the pids matching name"""
  75
  76     if mozinfo.isWin:
  77         # use the windows-specific implementation
  78         import wpk
  79         return wpk.get_pids(name)
  80     else:
  81         return [pid for pid,_ in running_processes(name)]
  82
  83 if __name__ == '__main__':
  84     pids = set()
  85     for i in sys.argv[1:]:
  86         pids.update(get_pids(i))
  87     for i in sorted(pids):
  88         print i