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

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

4 files changed, 73 insertions, 39 deletions

diff --git a/bitbake/bin/bitbake b/bitbake/bin/bitbake
index b046156b49..797b5a8d60 100755
--- a/bitbake/bin/bitbake
+++ b/bitbake/bin/bitbake
@@ -180,7 +180,7 @@ Default BBFILES are the .bb files in the current directory.""")
 
     serverinfo = server.BitbakeServerInfo(cooker.server)
 
-    server.BitBakeServerFork(serverinfo, cooker.serve, cooker_logfile)
+    server.BitBakeServerFork(cooker, cooker.server, serverinfo, cooker_logfile)
     del cooker
 
     # Setup a connection to the server (cooker)
@@ -203,7 +203,7 @@ Default BBFILES are the .bb files in the current directory.""")
         print("Valid interfaces are 'ncurses', 'depexp' or the default, 'knotty'.")
     else:
         try:
-            return_value = ui_init(serverConnection.connection, serverConnection.events)
+            return_value = server.BitbakeUILauch().launch(serverinfo, ui_init, serverConnection.connection, serverConnection.events)
         except Exception as e:
             print("FATAL: Unable to start to '%s' UI: %s" % (ui, e))
             raise
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
diff --git a/bitbake/lib/bb/server/none.py b/bitbake/lib/bb/server/none.py
index 14915d9737..38f713c519 100644
--- a/bitbake/lib/bb/server/none.py
+++ b/bitbake/lib/bb/server/none.py
@@ -109,9 +109,11 @@ class BitBakeServer():
     # remove this when you're done with debugging
     # allow_reuse_address = True
 
-    def __init__(self, cooker):
+    def __init__(self, cooker, pre_serve, post_serve):
         self._idlefuns = {}
         self.commands = BitBakeServerCommands(self, cooker)
+        self.pre_serve = pre_serve
+        self.post_serve = post_serve
 
     def register_idle_function(self, function, data):
         """Register a function to be called while the server is idle"""
@@ -160,9 +162,17 @@ class BitbakeServerInfo():
         self.commands = server.commands
 
 class BitBakeServerFork():
-    def __init__(self, serverinfo, command, logfile):
-        serverinfo.forkCommand = command
+    def __init__(self, cooker, server, serverinfo, logfile):
         serverinfo.logfile = logfile
+        serverinfo.cooker = cooker
+        serverinfo.server = server
+
+class BitbakeUILauch():
+    def launch(self, serverinfo, uifunc, *args):
+        serverinfo.server.pre_serve()
+        ret = bb.cooker.server_main(serverinfo.cooker, uifunc, *args)
+        serverinfo.server.post_serve()
+        return ret
 
 class BitBakeServerConnection():
     def __init__(self, serverinfo):
diff --git a/bitbake/lib/bb/server/xmlrpc.py b/bitbake/lib/bb/server/xmlrpc.py
index cb2949fb9f..30b83d2b13 100644
--- a/bitbake/lib/bb/server/xmlrpc.py
+++ b/bitbake/lib/bb/server/xmlrpc.py
@@ -87,7 +87,7 @@ class BitBakeServer(SimpleXMLRPCServer):
     # remove this when you're done with debugging
     # allow_reuse_address = True
 
-    def __init__(self, cooker, interface = ("localhost", 0)):
+    def __init__(self, cooker, pre_serve, post_serve, interface = ("localhost", 0)):
         """
         Constructor
         """
@@ -99,6 +99,9 @@ class BitBakeServer(SimpleXMLRPCServer):
         #self.register_introspection_functions()
         commands = BitBakeServerCommands(self, cooker)
         self.autoregister_all_functions(commands, "")
+        self.cooker = cooker
+        self.pre_serve = pre_serve
+        self.post_serve = post_serve
 
     def autoregister_all_functions(self, context, prefix):
         """
@@ -116,9 +119,14 @@ class BitBakeServer(SimpleXMLRPCServer):
         self._idlefuns[function] = data
 
     def serve_forever(self):
+        bb.cooker.server_main(self.cooker, self._serve_forever)
+
+    def _serve_forever(self):
         """
         Serve Requests. Overloaded to honor a quit command
         """
+        self.pre_serve()
+
         self.quit = False
         self.timeout = 0 # Run Idle calls for our first callback
         while not self.quit:
@@ -155,6 +163,7 @@ class BitBakeServer(SimpleXMLRPCServer):
             except:
                 pass
 
+        self.post_serve()
         self.server_close()
         return
 
@@ -164,8 +173,12 @@ class BitbakeServerInfo():
         self.port = server.port
 
 class BitBakeServerFork():
-    def __init__(self, serverinfo, command, logfile):
-        daemonize.createDaemon(command, logfile)
+    def __init__(self, cooker, server, serverinfo, logfile):
+        daemonize.createDaemon(server.serve_forever, logfile)
+
+class BitbakeUILauch():
+    def launch(self, serverinfo, uifunc, *args):
+        return uifunc(*args)
 
 class BitBakeServerConnection():
     def __init__(self, serverinfo):