From 4e3c768b0dae4895d942e71f7220fbfd3bc0d9f1 Mon Sep 17 00:00:00 2001
From: Steve Wattam <stephenwattam@gmail.com>
Date: Thu, 10 Oct 2013 11:32:37 +0100
Subject: [PATCH] Initial commit of edits made for tmatrix

---
 README.md | 20 +++++++++++---
 memusg    | 80 ++++++++++++++++++++++++++++++++++++++++++++-----------
 2 files changed, 81 insertions(+), 19 deletions(-)

diff --git a/README.md b/README.md
index 16f76f0..911ec3a 100644
--- a/README.md
+++ b/README.md
@@ -1,12 +1,24 @@
-A 'time'-like utility for Unix that measures peak memory usage.
+About
+-----
+Memusg is a `time`-like utility for Unix that measures peak memory usage.  It works by repeatedly polling using `ps`.
 
-Works in both interactive and non-interactive environments.
+Features
+--------
 
-Usage:
+ * Periodic memory checking
+ * Output to file or inline (`stderr`)
+ * Configurable interval
+ * Works in both interactive and non-interactive environments
+ * Passes signals through to the client application
+
+
+
+Usage
+-----
 
 ```bash
 export PATH=$path_to_memusg
-memusg my_command
+memusg [OPTS] my_command
 ```
 
 Example:
diff --git a/memusg b/memusg
index 07e1cb5..05e5686 100755
--- a/memusg
+++ b/memusg
@@ -1,4 +1,8 @@
-#!/usr/bin/env python
+#!/bin/sh
+''''exec python -u -- "$0" ${1+"$@"} # '''
+# vi: syntax=python
+
+import signal
 import sys
 import os
 from time import sleep
@@ -6,11 +10,27 @@ from subprocess import *
 
 # -o file.txt : output information to a file rather than stderr
 # -d          : debug mode
+# -m          : monitor mode: output every measurement
+# -p rate     : poll every rate seconds (default is 0.1)
+# -h          : Show usage info
+
+out         = sys.stderr
+DEBUG       = False
+poll_rate   = 0.1
+monitor     = False
+SIGNALS     = [signal.SIGINT, signal.SIGTERM, signal.SIGABRT]
 
+def print_usage(cmd):
+    print("USAGE: {} [-d] [-m] [-p TIME] [-o FILE] COMMAND [ARG [ARG [...]]]".format(cmd))
+    print("")
+    print(" -m      : Output every tick instead of at the end")
+    print(" -p TIME : Set the time period in seconds")
+    print(" -o FILE : Output to file (default is stdout)")
+    print(" -d      : Debug mode (verbose output)")
+    print("")
 
-out = sys.stderr
-DEBUG = False
 
+# Handle command-line input
 child_args = []
 i = 1
 while i < len(sys.argv):
@@ -19,6 +39,14 @@ while i < len(sys.argv):
         out = open(sys.argv[i], 'w')
     elif sys.argv[i] == '-d':
         DEBUG = True
+    elif sys.argv[i] == '-m':
+        monitor = True
+    elif sys.argv[i] == '-p':
+        i += 1
+        poll_rate = float(sys.argv[i])
+    elif sys.argv[i] == '-h':
+        print_usage(sys.argv[0])
+        sys.exit(0)
     else:
         child_args.append(sys.argv[i])
     i += 1
@@ -31,6 +59,10 @@ def log(msg):
     if DEBUG:
         print >>sys.stderr, "memusg: {}".format(msg)
 
+def put(msg):
+    out.write(msg)
+    out.flush # Only works if run with python -u
+
 def get_vsize(sid):
     vsize = 0
     # Example: /bin/ps -o vsize= --sid 23928
@@ -42,27 +74,42 @@ def get_vsize(sid):
         vsize += int(line.strip())
     return vsize
 
+def handle_signal(signal, frame):
+    # If the proc is running, pass the signal down.
+    if proc and not proc.returncode:
+        proc.send_signal(signal)
+    if fork_pid != 0:
+        os.kill(fork_pid, signal)
+    else:
+        sys.exit(0)
+
+
 # Create a new process session for this process so that we can
 # easily calculate the memory usage of the whole process tree using ps
 #
 # Since we need a new session using os.setsid(), we must first fork()
-pid = os.getpid()
-sid = os.getsid(pid)
-pgid = os.getpgid(pid)
+pid     = os.getpid()
+sid     = os.getsid(pid)
+pgid    = os.getpgid(pid)
 log("Pre-fork: PID is {} ; PGID is {} ; SID is {}".format(pid, pgid, sid))
 
-fork_pid = os.fork()
+fork_pid    = os.fork()
+proc        = None
+# Attach signal handler to everything we want to pass through
+for s in SIGNALS:
+    signal.signal(s, handle_signal)
+
 if fork_pid == 0:
     # We *are* the new fork (not the original process)
-    pid = os.getpid()
-    sid = os.getsid(pid)
-    pgid = os.getpgid(pid)
+    pid     = os.getpid()
+    sid     = os.getsid(pid)
+    pgid    = os.getpgid(pid)
     log("Post-fork: PID is {} ; PGID is {} ; SID is {}".format(pid, pgid, sid))
 
     log("Trying to init our own session".format(pid, pgid))
     os.setsid()
-    sid = os.getsid(pid)
-    pgid = os.getpgid(pid)
+    sid     = os.getsid(pid)
+    pgid    = os.getpgid(pid)
     log("Post-session init: PID is {} ; PGID is {} ; SID is {}".format(pid, pgid, sid))
         
     log("Starting child: {}".format(child_command))
@@ -71,12 +118,15 @@ if fork_pid == 0:
 
     vmpeak = -1
     while proc.returncode == None:
-        vmpeak = max(get_vsize(sid), vmpeak)
+        size = get_vsize(sid)
+        if monitor:
+            put("memusg: {} kb\n".format(size))
+        vmpeak = max(size, vmpeak)
         log("Waiting for child to exit. vmpeak={}".format(vmpeak))
         proc.poll()
-        sleep(0.1) # Time in seconds (float)
+        sleep(poll_rate) # Time in seconds (float)
     
-    out.write("memusg: vmpeak: {} kb\n".format(vmpeak))
+    put("memusg: vmpeak: {} kb\n".format(vmpeak))
 
     status = proc.returncode
     log("Child process returned {}".format(status))