Add profiling support.

svn path=/trunk/; revision=4636

1 files changed, 28 insertions, 12 deletions

diff --git a/rpkid/rpki/rpkic.py b/rpkid/rpki/rpkic.py
index 9bcf62a1..5f8ff31a 100644
--- a/rpkid/rpki/rpkic.py
+++ b/rpkid/rpki/rpkic.py
@@ -81,23 +81,39 @@ class main(rpki.cli.Cmd):
 
     rpki.log.use_syslog = False
 
-    cfg_file = None
-    handle = None
+    self.cfg_file = None
+    self.handle = None
+    profile = None
 
-    opts, argv = getopt.getopt(sys.argv[1:], "c:hi:?", ["config=", "help", "identity="])
+    opts, self.argv = getopt.getopt(sys.argv[1:], "c:hi:?",
+                                    ["config=", "help", "identity=", "profile="])
     for o, a in opts:
       if o in ("-c", "--config"):
-        cfg_file = a
+        self.cfg_file = a
       elif o in ("-h", "--help", "-?"):
-        argv = ["help"]
+        self.argv = ["help"]
       elif o in ("-i", "--identity"):
-        handle = a
-
-    if not argv or argv[0] != "help":
-      rpki.log.init("rpkic")
-      self.read_config(cfg_file, handle)
-
-    rpki.cli.Cmd.__init__(self, argv)
+        self.handle = a
+      elif o == "--profile":
+        profile = a
+
+    if self.argv and self.argv[0] == "help":
+      rpki.cli.Cmd.__init__(self, self.argv)
+    elif profile:
+      import cProfile
+      prof = cProfile.Profile()
+      try:
+        prof.runcall(self.main)
+      finally:
+        prof.dump_stats(profile)
+        print "Dumped profile data to %s" % profile
+    else:
+      self.main()
+
+  def main(self):
+    rpki.log.init("rpkic")
+    self.read_config(self.cfg_file, self.handle)
+    rpki.cli.Cmd.__init__(self, self.argv)
 
   def read_config(self, cfg_file, handle):
     global rpki