1 import cProfile

     2 import ppstats

     3 

     4 from google.appengine.ext import webapp

     5 from google.appengine.api import memcache

     6 import google.appengine.ext.webapp.util

     7 

     8 import os.path

     9 import logging

    10 import re

    11 import random

    12 import string

    13 import zlib

    14 

    15 mc_client = memcache.Client()

    16 

    17 alphanumeric = string.letters + string.digits

    18 

    19 global_profiler = None

    20 

    21 class GAEProfiler(object):

    22     _save_every = 10

    23 

    24     def __init__(self):

    25         self.is_profiling = False

    26         self._profiler = None

    27         self.num_requests = 0

    28         self.requests_profiled = 0

    29         self.request_regex = None

    30         self.profile_key = ''.join([random.choice(alphanumeric) for x in range(4)])

    31 

    32     def start_profiling(self, request_regex=None, num_requests=0):

    33         "start profiling with this object, setting # of requests and filter"

    34         if self.is_profiling:

    35             return

    36 

    37         self.is_profiling = True

    38         if self._profiler is None:

    39             self._profiler = cProfile.Profile()

    40         self.num_requests = num_requests

    41         if request_regex:

    42             self.request_regex = re.compile(request_regex)

    43 

    44     def stop_profiling(self):

    45         self.is_profiling = False

    46 

    47     def resume_profiling(self):

    48         self.is_profiling = True

    49     

    50     def has_profiler(self):

    51         return self._profiler is not None

    52 

    53     def get_pstats(self):

    54         "return a ppstats object from current profile data"

    55         gae_base_dir = '/'.join(webapp.__file__.split('/')[:-5])

    56         sys_base_dir = '/'.join(logging.__file__.split('/')[:-2])

    57 

    58         stats = ppstats.Stats(self._profiler)

    59         stats.hide_directory(gae_base_dir, 'GAEHome')

    60         stats.hide_directory(sys_base_dir, 'SysHome')

    61         stats.strip_dirs()

    62         return stats

    63 

    64     def runcall(self, func, *args, **kwargs):

    65         "profile one call, incrementing requests_profiled and maybe saving stats"

    66         self.requests_profiled += 1

    67         if self._profiler:

    68             ret = self._profiler.runcall(func, *args, **kwargs)

    69         else:

    70             ret = func(*args, **kwargs)

    71         

    72 #        if (self.requests_profiled % self._save_every) == 0 or \

    73 #                self.requests_profiled == self.num_requests:

    74 #            self.save_pstats_to_memcache()

    75         self.save_pstats_to_memcache()

    76         return ret

    77 

    78     def should_profile_request(self):

    79         "check for # of requests profiled and that SCRIPT_NAME matches regex"

    80         env = dict(os.environ)

    81         script_name = env.get('SCRIPT_NAME', '')

    82         logging.info(script_name)

    83         

    84         if self.num_requests and self.requests_profiled >= self.num_requests:

    85             return False

    86 

    87         if self.request_regex and not self.request_regex.search(script_name):

    88             return False

    89 

    90         return True

    91 

    92     def save_pstats_to_memcache(self):

    93         "save stats from profiler object to memcache"

    94         ps = self.get_pstats()

    95         output = ps.dump_stats_pickle()

    96         compressed_data = zlib.compress(output, 3)

    97         cache_key = cache_key_for_profile(self.profile_key)

    98         mc_client.set(cache_key, compressed_data)

    99         logging.info("Saved pstats to memcache with key %s" % cache_key)

   100 

   101 

   102 

   103 def get_global_profiler():

   104     global global_profiler

   105     if not global_profiler:

   106         global_profiler = GAEProfiler()

   107 

   108     return global_profiler

   109     

   110 def cache_key_for_profile(profile_key):

   111     "generate a memcache key"

   112     return "ProfileData.%s" % profile_key

   113 

   114 def load_pstats_from_memcache(profile_key):

   115     "retrieve ppstats object"

   116     mc_data = mc_client.get(cache_key_for_profile(profile_key))

   117     if not mc_data:

   118         return None

   119 

   120     return ppstats.from_gz(mc_data)

   121 

   122 def get_stats_from_global_or_request(request_obj):

   123     "get pstats for a key, or the global pstats"

   124     key = request_obj.get('key', '')

   125     if key:

   126         return load_pstats_from_memcache(key)

   127     else:

   128         gp = get_global_profiler()

   129         if not gp.has_profiler():

   130             return None

   131         return gp.get_pstats()

   132 

   133 class show_profile(webapp.RequestHandler):

   134     def get(self):

   135         ps = get_stats_from_global_or_request(self.request)

   136         if not ps:

   137             self.response.out.write("<body><html><h3>No profiler.</h3><html></body>")

   138             return

   139 

   140         ps.set_output(self.response.out)

   141         sort = self.request.get('sort', 'time')

   142         ps.sort_stats(sort)

   143         self.response.out.write("<body><html><pre>\n")

   144         ps.print_stats(30)

   145         self.response.out.write("</pre></html></body>")

   146 

   147 class download_profile_data(webapp.RequestHandler):

   148     def get(self):

   149         ps = get_stats_from_global_or_request(self.request)

   150         if not ps:

   151             self.response.out.write("<body><html><h3>No profiler.</h3><html></body>")

   152             return            

   153 

   154         output = ps.dump_stats_pickle()

   155 

   156         self.response.headers['Content-Type'] = 'application/octet-stream'

   157 

   158         self.response.out.write(output)

   159 

   160 class show_profiler_status(webapp.RequestHandler):

   161     def get(self):

   162         gp = get_global_profiler()

   163         if not gp.has_profiler:

   164             self.response.out.write("<body><html><h3>No profiler.</h3><html></body>")

   165             return

   166 

   167         self.response.out.write("<html><body>")

   168         self.response.out.write("<b>Currently profiling:</b> %s<br>" % gp.is_profiling)

   169         self.response.out.write("<b>Profile Key</b>: %s<br>" % gp.profile_key)

   170         self.response.out.write("<b>Requests profiled so far:</b> %s<br>" % gp.requests_profiled)

   171         self.response.out.write("<b>Requests to profile:</b> %s<br>" % gp.num_requests)

   172         self.response.out.write("<b>Request regex:</b> %s<br>" % gp.request_regex)

   173         self.response.out.write("</body></html>")

   174 

   175 class start_profiler(webapp.RequestHandler):

   176     def get(self):

   177         gp = get_global_profiler()

   178         gp.start_profiling()

   179         self.response.headers['Content-Type'] = "text/plain"

   180         self.response.out.write("Started profiling (key: %s).\n" % gp.profile_key)

   181         self.response.out.write("Retrieve saved results at <a href='/profiler/show?key=%(key)s'>/profiler/show?key=%(key)s).\n" % {'key':gp.profile_key})

   182 

   183 class stop_profiler(webapp.RequestHandler):

   184     def get(self):

   185         gp = get_global_profiler()

   186         gp.stop_profiling()

   187         self.request.out.write("Content-Type: text/plain\n\n")

   188         self.request.out.write("done.")

   189 

   190 class save_profile_data(webapp.RequestHandler):

   191     def get(self):

   192         gp = get_global_profiler()

   193         

   194 

   195 def _add_our_endpoints(application):

   196     "insert our URLs into the application map"

   197     url_mapping = [(regex.pattern, handler) for (regex, handler) in application._url_mapping]

   198     return webapp.WSGIApplication(url_mapping, debug=True)

   199 

   200 #

   201 #  wrapper to for webapp applications

   202 #

   203 def run_wsgi_app(application):

   204     "proxy webapp.util's call to profile when needed"

   205     gp = get_global_profiler()

   206     if gp.is_profiling and gp.should_profile_request():

   207         return gp.runcall(google.appengine.ext.webapp.util.run_wsgi_app, *(application,))

   208     else:

   209         return google.appengine.ext.webapp.util.run_wsgi_app(application)

   210 

   211 #

   212 # middleware for django applications

   213 #

   214 

   215 class ProfileMiddleware(object):

   216     def __init__(self):

   217         self.profiler = None

   218 

   219     def process_request(self, request):

   220         self.profiler = get_global_profiler()

   221 

   222     def process_view(self, request, callback, callback_args, callback_kwargs):

   223         if self.profiler.is_profiling:

   224             return self.profiler.runcall(callback, request, *callback_args, **callback_kwargs)

   225 

   226 application = webapp.WSGIApplication(

   227     [('/profiler/start', start_profiler),

   228      ('/profiler/stop', stop_profiler),

   229      ('/profiler/show', show_profile),

   230      ('/profiler/download', download_profile_data),

   231      ('/profiler/status', show_profiler_status),

   232      ],

   233     debug=True)

   234 

   235     

   236 def main():

   237     google.appengine.ext.webapp.util.run_wsgi_app(application)

   238 

   239 if __name__ == '__main__':

   240     main()