Source

hotdebug / hook.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# vim: set ai et si ts=4 sts=4 sw=4:

from __future__ import with_statement
import sys, threading
from cProfile import Profile
from cStringIO import StringIO
try:
    import pstats
except ImportError:
    pstats = None


class RLock(object):

    def __init__(self):
        self._lock = threading.RLock()

    def __enter__(self):
        self._lock.acquire()
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        self._lock.release()


def _hook_export_service():

    def public(fn):
        globals()[fn.__name__] = fn
        return fn

    _wrappers = []
    def wrapper(fn):
        _wrappers.append(fn)
        return fn

    # ===  catch exceptions ===

    # exception callback lock held when processing callbacks
    _exc_lock = RLock()

    _exc_cbs = {}

    @public
    def register_exception_callback(id, cb):
        '''
        Register callback that will be called when has uncaught exception
        '''
        with _exc_lock:
            _exc_cbs[id] = cb

    @public
    def unregister_exception_callback(id):
        '''
        Reverse op of register_exception_callback
        '''
        with _exc_lock:
            _exc_cbs.pop(id, None)

    @wrapper
    def wrap_catch_exc(fn):

        def wrapped(*args, **kw):
            try:
                return fn(*args, **kw)
            except:
                # there may be several threads that raise exception
                # so one-by-one
                last_type, last_value, last_traceback = sys.exc_info()
                with _exc_lock:
                    for cb in _exc_cbs.itervalues():
                        cb(last_type, last_value, last_traceback)
                raise

        return wrapped

    # === profile ===

    _profile_objs = {}

    _profile_lock = RLock()

    @public
    def start_profile(id):
        '''
        Start to profile
        '''
        if pstats is None:
            return False

        if _profile_objs:
            return False
        with _profile_lock:
            _profile_objs[id] = Profile()
        return True
    
    @public
    def stop_profile(id, sort="time"):
        '''
        Stop to profile and return result
        '''
        if pstats is None:
            return None

        with _profile_lock:
            p = _profile_objs.pop(id, None)
            if p is None:
                return None
            result = StringIO()
            pstats.Stats(p, stream=result).strip_dirs().sort_stats(
                    sort).print_stats()
            return result.getvalue()

    @wrapper
    def wrap_profile(fn):

        def wrapped(*args, **kw):
            if _profile_objs:
                with _profile_lock:
                    _profile_obj = _profile_objs.values()[0]
                    return _profile_obj.runcall(fn, *args, **kw)
            return fn(*args, **kw)
        
        return wrapped

    # === modify ===

    from netsvc import ExportService
    for service in ExportService._services.itervalues():
        dispatch = service.__class__.__dict__.get('dispatch', None)
        if dispatch and not getattr(dispatch, '_wrapped', False):
            for w in _wrappers:
                dispatch = w(dispatch)
            dispatch._wrapped = True
            service.__class__.dispatch = dispatch


_hook_export_service()
del _hook_export_service