bitbake: Rewrite profiling code so its functional for both none and xmlrpc backends

Signed-off-by: Richard Purdie <rpurdie@linux.intel.com>

1 files changed, 42 insertions, 31 deletions

diff --git a/bitbake/lib/bb/cooker.py b/bitbake/lib/bb/cooker.py
index 44b9b2c31a..33eb65e2f3 100644
--- a/bitbake/lib/bb/cooker.py
+++ b/bitbake/lib/bb/cooker.py
@@ -1,4 +1,3 @@
-#!/usr/bin/env python
 # ex:ts=4:sw=4:sts=4:et
 # -*- tab-width: 4; c-basic-offset: 4; indent-tabs-mode: nil -*-
 #
@@ -71,7 +70,7 @@ class BBCooker:
         self.bb_cache = None
 
         if server:
-            self.server = server.BitBakeServer(self)
+            self.server = server.BitBakeServer(self, self.pre_serve, self.post_serve)
 
         self.configuration = configuration
 
@@ -916,41 +915,53 @@ class BBCooker:
             return self.appendlist[f] 
         return []
 
-    def serve(self):
-
+    def pre_serve(self):
         # Empty the environment. The environment will be populated as
         # necessary from the data store.
         bb.utils.empty_environment()
 
-        if self.configuration.profile:
-            try:
-                import cProfile as profile
-            except:
-                import profile
-
-            profile.runctx("self.server.serve_forever()", globals(), locals(), "profile.log")
-
-            # Redirect stdout to capture profile information
-            pout = open('profile.log.processed', 'w')
-            so = sys.stdout.fileno()
-            os.dup2(pout.fileno(), so)
-
-            import pstats
-            p = pstats.Stats('profile.log')
-            p.sort_stats('time')
-            p.print_stats()
-            p.print_callers()
-            p.sort_stats('cumulative')
-            p.print_stats()
-
-            os.dup2(so, pout.fileno())
-            pout.flush()
-            pout.close()
-        else:
-            self.server.serve_forever()
-
+    def post_serve(self):
         bb.event.fire(CookerExit(), self.configuration.event_data)
 
+
+def server_main(cooker, func, *args):
+    if cooker.configuration.profile:
+        try:
+            import cProfile as profile
+        except:
+            import profile
+        prof = profile.Profile()
+
+        ret = profile.Profile.runcall(prof, func, *args)
+
+        prof.dump_stats("profile.log")
+
+        # Redirect stdout to capture profile information
+        pout = open('profile.log.processed', 'w')
+        so = sys.stdout.fileno()
+        orig_so = os.dup(sys.stdout.fileno())
+        os.dup2(pout.fileno(), so)
+   
+        import pstats
+        p = pstats.Stats('profile.log')
+        p.sort_stats('time')
+        p.print_stats()
+        p.print_callers()
+        p.sort_stats('cumulative')
+        p.print_stats()
+
+        os.dup2(orig_so, so)
+        pout.flush()
+        pout.close()  
+
+        print("Raw profiling information saved to profile.log and processed statistics to profile.log.processed")
+
+        return ret
+    else:
+        return func(*args)
+
+
+
 class CookerExit(bb.event.Event):
     """
     Notify clients of the Cooker shutdown