tools/etc/memusg.mac

   1 #! /bin/sh
   2 """:"
   3 exec python $0 ${1+"$@"}
   4 """
   5 # vi: syntax=python
   6
   7 import signal
   8 import sys
   9 import os
  10 from time import sleep
  11 from subprocess import *
  12
  13 # -o file.txt : output information to a file rather than stderr
  14 # -d          : debug mode
  15 # -m          : monitor mode: output every measurement
  16 # -p rate     : poll every rate seconds (default is 0.1)
  17 # -h          : Show usage info
  18
  19 out = sys.stderr
  20 DEBUG = False
  21 poll_rate = 0.1
  22 monitor = False
  23 SIGNALS = [signal.SIGINT, signal.SIGTERM, signal.SIGABRT]
  24
  25 def print_usage(cmd):
  26     print("USAGE: {} [-d] [-m] [-p TIME] [-o FILE] COMMAND [ARG [ARG [...]]]".format(cmd))
  27     print("")
  28     print(" -m      : Output every tick instead of at the end")
  29     print(" -p TIME : Set the time period in seconds")
  30     print(" -o FILE : Output to file (default is stdout)")
  31     print(" -d      : Debug mode (verbose output)")
  32     print("")
  33
  34 # Handle command-line input
  35 child_args = []
  36 i = 1
  37 while i < len(sys.argv):
  38     if sys.argv[i] == '-o':
  39         i += 1
  40         out = open(sys.argv[i], 'w')
  41     elif sys.argv[i] == '-d':
  42         DEBUG = True
  43     elif sys.argv[i] == '-m':
  44         monitor = True
  45     elif sys.argv[i] == '-p':
  46         i += 1
  47         poll_rate = float(sys.argv[i])
  48     elif sys.argv[i] == '-h':
  49         print_usage(sys.argv[0])
  50         sys.exit(0)
  51     else:
  52         child_args.append(sys.argv[i])
  53     i += 1
  54
  55 # child_command should be a single argument as if to "/bin/sh -c 'child_command'"
  56 # when shell=True is enabled
  57 child_command = ' '.join(child_args)
  58
  59 def log(msg):
  60     if DEBUG:
  61         print >>sys.stderr, "memusg: {}".format(msg)
  62 def put(msg):
  63     out.write(msg)
  64     out.flush # Only works if run with python -u
  65
  66 def get_vsize(sid):
  67     vsize = 0
  68     # Example: /bin/ps -o vsize= --sid 23928
  69     proc = Popen(['ps', '-o', 'vsize=', '-g', str(sid)], stdout=PIPE, stderr=None, shell=False)
  70     (stdout, _stderr) = proc.communicate()
  71     # Iterate over each process within the process tree of our process session
  72     # (this ensures that we include processes launched by a child bash script, etc.)
  73     for line in stdout.split():
  74         vsize += int(line.strip())
  75     return vsize
  76
  77 def handle_signal(signal, frame):
  78     # If the proc is running, pass the signal down.
  79     if proc and not proc.returncode:
  80         proc.send_signal(signal)
  81     if fork_pid != 0:
  82         os.kill(fork_pid, signal)
  83     else:
  84         sys.exit(0)
  85
  86 # Create a new process session for this process so that we can
  87 # easily calculate the memory usage of the whole process tree using ps
  88 #
  89 # Since we need a new session using os.setsid(), we must first fork()
  90 pid = os.getpid()
  91 sid = os.getsid(pid)
  92 pgid = os.getpgid(pid)
  93 log("Pre-fork: PID is {} ; PGID is {} ; SID is {}".format(pid, pgid, sid))
  94
  95 fork_pid = os.fork()
  96
  97 # Attach signal handler to everything we want to pass through
  98 proc = None   # after ^C, returns here and need `proc` being defined.
  99 for s in SIGNALS:
 100     signal.signal(s, handle_signal)
 101
 102 if fork_pid == 0:
 103     # We *are* the new fork (not the original process)
 104     pid = os.getpid()
 105     sid = os.getsid(pid)
 106     pgid = os.getpgid(pid)
 107     log("Post-fork: PID is {} ; PGID is {} ; SID is {}".format(pid, pgid, sid))
 108
 109     log("Trying to init our own session".format(pid, pgid))
 110     os.setsid()
 111     sid = os.getsid(pid)
 112     pgid = os.getpgid(pid)
 113     log("Post-session init: PID is {} ; PGID is {} ; SID is {}".format(pid, pgid, sid))
 114
 115     log("Starting child: {}".format(child_command))
 116     # "None" means "inherit from parent"
 117     proc = Popen(child_command, stdin=None, stdout=None, stderr=None, env=None, shell=True)
 118
 119     vmpeak = -1
 120     while proc.returncode == None:
 121         #vmpeak = max(get_vsize(pgid), vmpeak)
 122         size = get_vsize(pgid)
 123         if monitor:
 124             put("memusg: {} kb\n".format(size))
 125         vmpeak = max(size, vmpeak)
 126         log("Waiting for child to exit. vmpeak={}".format(vmpeak))
 127         proc.poll()
 128         sleep(poll_rate) # Time in seconds (float)
 129
 130     put("memusg: vmpeak: {} kb\n".format(vmpeak))
 131
 132     status = proc.returncode
 133     log("Child process returned {}".format(status))
 134     sys.exit(status)
 135
 136 else:
 137     # This is the branch of fork that continues the original process
 138     (_fork_pid, full_status) = os.waitpid(fork_pid, 0)
 139     status = full_status >> 8
 140     log("Fork returned {}".format(status))
 141     out.close()
 142     sys.exit(status)