Add option to print stacktrace of all active threads

lptrace

@@ -7,9 +7,10 @@ import signal
 import tempfile
 import subprocess
 
-from optparse import OptionParser
+from argparse import ArgumentParser
 
-trace_fn = """def __lptrace_trace_calls__(frame, event, arg):
+trace_fn = """
+def __lptrace_trace_calls__(frame, event, arg):
     if event != 'call':
         return
 
@@ -27,22 +28,30 @@ import sys ; sys.settrace(__lptrace_trace_calls__)"""
 untrace_fn = """import sys ; sys.settrace(None)"""
 
 
-def runfile(pid, script):
-    with tempfile.NamedTemporaryFile() as tmp:
-        name = tmp.name
-        tmp.write(script)
-
-        tmp.file.flush()
-        os.chmod(tmp.name, 0666)
+def runfile(pid, script, autoclose=True):
+    tmp = tempfile.NamedTemporaryFile();
+    name = tmp.name
+    tmp.write(script)
+    tmp.file.flush()
+    os.chmod(tmp.name, 0666)
+    cmd = 'execfile(\\"{}\\")'.format(name)
+    p = inject(pid, cmd)
+    if autoclose:
+        p.wait()
+        tmp.close()
+        return None
+    else:
+        return p, tmp
 
-        cmd = 'execfile(\\"{}\\")'.format(name)
-        inject(pid, cmd)
 
+def create_fifo():
+    fifo = tempfile.mktemp()
+    os.mkfifo(fifo)
+    os.chmod(fifo, 0777)
+    return fifo
 
 def strace(pid):
-    fifo_name = tempfile.mktemp()
-    os.mkfifo(fifo_name)
-    os.chmod(fifo_name, 0777)
+    fifo_name = create_fifo()
 
     trace_code = trace_fn % fifo_name
     runfile(pid, trace_code)
@@ -62,6 +71,50 @@ def strace(pid):
                 print data
 
 
+print_stacktrace_fn = """
+import traceback
+import threading
+import sys
+
+def print_stacktrace(out_fifo):
+    out_channel = open(out_fifo, 'w')
+    #traceback.print_stack(file=out_channel)
+    print len(threading.enumerate())
+    for th in threading.enumerate():
+        out_channel.write('%%r\\n' %% th)
+        traceback.print_stack(sys._current_frames()[th.ident], file=out_channel)
+        out_channel.write('\\n')
+    out_channel.close()
+
+OUTPUT_FIFO = '%s'
+
+print_stacktrace(OUTPUT_FIFO)
+"""
+
+
+def print_stacktrace(pid):
+    def sigint_handler(signal, frame):
+        print 'Received Ctrl-C, quitting'
+        sys.exit(0)
+
+    signal.signal(signal.SIGINT, sigint_handler)
+
+    out_fifo = create_fifo()
+    code = print_stacktrace_fn % (out_fifo)
+
+    p, tmp_f = runfile(pid, code, autoclose=False)
+
+    with open(out_fifo, 'r') as of:
+        while True:
+            l = of.read()
+            if l == '':
+                break
+            print l
+
+    p.wait()
+    tmp_f.close()
+
+
 def pdb_prompt(pid):
     code = 'import pdb ; pdb.set_trace()'
     inject(pid, code)
@@ -78,9 +131,11 @@ def inject(pid, code):
     cmdline = 'gdb -p %d -batch %s' % (pid,
         ' '.join(["-eval-command='call %s'" % cmd for cmd in gdb_cmds]))
 
+    #print 'Executing %s' % cmdline
+
     p = subprocess.Popen(cmdline, shell=True, stdout=subprocess.PIPE,
                          stderr=subprocess.PIPE, )
-    out, err = p.communicate()
+    return p
 
 
 def cmd_exists(cmd):
@@ -89,21 +144,22 @@ def cmd_exists(cmd):
 
 
 def main():
-    parser = OptionParser()
-    parser.add_option("-p", "--process", dest="pid",
-                      help="Attach to prod $pid")
+    parser = ArgumentParser('lptrace')
+    parser.add_argument("-p", "--process", dest="pid", required=True,
+                        help="Attach to prod $pid")
 
-    parser.add_option("-d", "--debugger",
-                      action="store_true", dest="debugger", default=False,
-                      help="Inject a pdb prompt")
+    group = parser.add_mutually_exclusive_group()
 
-    (options, args) = parser.parse_args()
+    group.add_argument("-d", "--debugger",
+                     action="store_true", dest="debugger", default=False,
+                     help="Inject a pdb prompt")
+    group.add_argument("-s", "--stacktrace",
+                     action="store_true", dest="stacktrace", default=False,
+                     help="Print stacktrace")
 
-    if options.pid is None:
-        print "You need to specify a process to attach to."
-        sys.exit(-1)
+    args = parser.parse_args()
 
-    pid = int(options.pid)
+    pid = int(args.pid)
 
     if not cmd_exists('gdb'):
         print "Error: lptrace requires GDB >= 7.x. Exiting."
@@ -113,8 +169,10 @@ def main():
         print "Error: you must be root to run lptrace. Exiting."
         sys.exit(-1)
 
-    if options.debugger is True:
+    if args.debugger:
         pdb_prompt(pid)
+    elif args.stacktrace:
+        print_stacktrace(pid)
     else:
         strace(pid)