cooker/process.py: Move profiling code to a place it can be reused by different server mechanisms

The cooker profiling code isn't server specific so move it to a place
where different server backends can use it.

Signed-off-by: Richard Purdie <richard.purdie@linuxfoundation.org>

2 files changed, 51 insertions, 27 deletions

diff --git a/lib/bb/cooker.py b/lib/bb/cooker.py
index 622227393..6096e54b0 100644
--- a/lib/bb/cooker.py
+++ b/lib/bb/cooker.py
@@ -1058,12 +1058,62 @@ class BBCooker:
             return self.appendlist[f]
         return []
 
+    def pre_serve(self):
+        # Empty the environment. The environment will be populated as
+        # necessary from the data store.
+        #bb.utils.empty_environment()
+        return
+
+    def post_serve(self):
+        bb.event.fire(CookerExit(), self.configuration.event_data)
+
     def shutdown(self):
         self.state = state.shutdown
 
     def stop(self):
         self.state = state.stop
 
+def server_main(cooker, func, *args):
+    cooker.pre_serve()
+
+    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")
+
+    else:
+        ret = func(*args)
+
+    cooker.post_serve()
+
+    return ret
+
 class CookerExit(bb.event.Event):
     """
     Notify clients of the Cooker shutdown
diff --git a/lib/bb/server/process.py b/lib/bb/server/process.py
index 5d7f8aa9d..88f819106 100644
--- a/lib/bb/server/process.py
+++ b/lib/bb/server/process.py
@@ -95,33 +95,7 @@ class ProcessServer(Process):
         self._idlefunctions[function] = data
 
     def run(self):
-        if self.configuration.profile:
-            return self.profile_main()
-        else:
-            return self.main()
-
-    def profile_main(self):
-        import cProfile
-        profiler = cProfile.Profile()
-        try:
-            return profiler.runcall(self.main)
-        finally:
-            profiler.dump_stats(self.profile_filename)
-            self.write_profile_stats()
-            sys.__stderr__.write("Raw profiling information saved to %s and "
-                                 "processed statistics to %s\n" %
-                                 (self.profile_filename,
-                                  self.profile_processed_filename))
-
-    def write_profile_stats(self):
-        import pstats
-        with open(self.profile_processed_filename, 'w') as outfile:
-            stats = pstats.Stats(self.profile_filename, stream=outfile)
-            stats.sort_stats('time')
-            stats.print_stats()
-            stats.print_callers()
-            stats.sort_stats('cumulative')
-            stats.print_stats()
+        bb.cooker.server_main(self.cooker, self.main)
 
     def main(self):
         # Ignore SIGINT within the server, as all SIGINT handling is done by