tools/etc/memusg.linux

   1 #!/usr/bin/env python
   2 import sys
   3 import os
   4 import signal
   5 from time import sleep
   6 from subprocess import *
   7
   8 # -o file.txt : output information to a file rather than stderr
   9 # -d          : debug mode
  10 # -m INT      : kill process when max kb reached
  11
  12
  13 out = sys.stderr
  14 DEBUG = False
  15 MAX_MEM = 0
  16
  17 child_args = []
  18 i = 1
  19 while i < len(sys.argv):
  20     if sys.argv[i] == '-o':
  21         i += 1
  22         out = open(sys.argv[i], 'w')
  23     elif sys.argv[i] == '-d':
  24         DEBUG = True
  25     elif sys.argv[i] == '-m':
  26         MAX_MEM = int(sys.argv[i+1])
  27         i += 1
  28     else:
  29         child_args.append(sys.argv[i])
  30     i += 1
  31
  32 # child_command should be a single argument as if to "/bin/sh -c 'child_command'"
  33 # when shell=True is enabled
  34 child_command = ' '.join(child_args)
  35
  36 def log(msg):
  37     if DEBUG:
  38         print >>sys.stderr, "memusg: {}".format(msg)
  39
  40 def get_vsize(sid):
  41     vsize = 0
  42     # Example: /bin/ps -o vsize= --sid 23928
  43     proc = Popen(['ps', '-o', 'vsize=', '--sid', str(sid)], stdout=PIPE, stderr=None, shell=False)
  44     (stdout, _stderr) = proc.communicate()
  45     # Iterate over each process within the process tree of our process session
  46     # (this ensures that we include processes launched by a child bash script, etc.)
  47     for line in stdout.split():
  48         vsize += int(line.strip())
  49     return vsize
  50
  51 # Create a new process session for this process so that we can
  52 # easily calculate the memory usage of the whole process tree using ps
  53 #
  54 # Since we need a new session using os.setsid(), we must first fork()
  55 pid = os.getpid()
  56 sid = os.getsid(pid)
  57 pgid = os.getpgid(pid)
  58 log("Pre-fork: PID is {} ; PGID is {} ; SID is {}".format(pid, pgid, sid))
  59
  60 fork_pid = os.fork()
  61 if fork_pid == 0:
  62     # We *are* the new fork (not the original process)
  63     pid = os.getpid()
  64     sid = os.getsid(pid)
  65     pgid = os.getpgid(pid)
  66     log("Post-fork: PID is {} ; PGID is {} ; SID is {}".format(pid, pgid, sid))
  67
  68     log("Trying to init our own session".format(pid, pgid))
  69     os.setsid()
  70     sid = os.getsid(pid)
  71     pgid = os.getpgid(pid)
  72     log("Post-session init: PID is {} ; PGID is {} ; SID is {}".format(pid, pgid, sid))
  73
  74     log("Starting child: {}".format(child_command))
  75     # "None" means "inherit from parent"
  76     proc = Popen(child_command, stdin=None, stdout=None, stderr=None, env=None, shell=True)
  77
  78     vmpeak = -1
  79     while proc.returncode == None:
  80         vmpeak = max(get_vsize(sid), vmpeak)
  81         if MAX_MEM > 0 and MAX_MEM < vmpeak:
  82             log("Max memory {} reached. Killing child.".format(MAX_MEM))
  83             os.killpg(int(pgid), signal.SIGTERM)
  84         log("Waiting for child to exit. vmpeak={}".format(vmpeak))
  85         proc.poll()
  86         sleep(0.1) # Time in seconds (float)
  87
  88     out.write("memusg: vmpeak: {} kb\n".format(vmpeak))
  89
  90     status = proc.returncode
  91     log("Child process returned {}".format(status))
  92     sys.exit(status)
  93
  94 else:
  95     # This is the branch of fork that continues the original process
  96     (_fork_pid, full_status) = os.waitpid(fork_pid, 0)
  97     status = full_status >> 8
  98     log("Fork returned {}".format(status))
  99     out.close()
 100     sys.exit(status)