Python tips: 逾時裝飾器, @timeout decorator

最近工作有點多，趁周末有空，繼續分享我在學習和使用python過程中的一些小tips。

有沒有遇到過這樣的事情：對資料庫執行插入或更新操作，因為資料量大或其他原因，導致此次操作非常耗時，有時甚至等上好幾個小時，也無法完成。很鬱悶，怎麼操作不逾時啊？因為資料庫配置時逾時時間很長，並且有些操作又是需要很長時間的，所以不能修改預設的逾時時間。

因為客觀條件不允許，我們不能靠資料庫逾時來終止此次操作，所以必須要在自己的方法邏輯模組裡實現逾時檢測的功能。

在python裡有沒有可以不用修改原來的方法內部邏輯，就能實現逾時檢測呢？肯定有啦，就是利用裝飾器。裝飾器是什嗎？在部落格園找到了一篇介紹文章：函數和方法裝飾漫談(Function decorator)。

廢話聽完，我現在介紹主角出場：逾時裝飾器，timeout decorator。

逾時檢測邏輯：啟動新子線程執行指定的方法，主線程等待子線程的運行結果，若在指定時間內子線程還未執行完畢，則判斷為逾時，拋出逾時異常，並殺掉子線程；否則未逾時，返回子線程所執行的方法的傳回值。

在實現過程中，發現python預設模組裡是沒有方法可以殺掉線程的，怎麼辦呢？當然先問問google或百度，果然，keill thread這個關鍵詞很熱門，很快就搜尋到我想要的東西了:"Kill a thread in Python"，就是以下這個KThread類，它繼承了threading.Thread，並添加了kill方法，讓我們能殺掉它：

 

import sys

class KThread(threading.Thread):
    """A subclass of threading.Thread, with a kill()
    method.
    
    Come from:
    Kill a thread in Python: 
    http://mail.python.org/pipermail/python-list/2004-May/260937.html
    """
    def __init__(self, *args, **kwargs):
        threading.Thread.__init__(self, *args, **kwargs)
        self.killed = False

    def start(self):
        """Start the thread."""
        self.__run_backup = self.run
        self.run = self.__run      # Force the Thread to install our trace.
        threading.Thread.start(self)

    def __run(self):
        """Hacked run function, which installs the
        trace."""
        sys.settrace(self.globaltrace)
        self.__run_backup()
        self.run = self.__run_backup

    def globaltrace(self, frame, why, arg):
        if why == 'call':
          return self.localtrace
        else:
          return None

    def localtrace(self, frame, why, arg):
        if self.killed:
          if why == 'line':
            raise SystemExit()
        return self.localtrace

    def kill(self):
        self.killed = True

 

好了，萬事戒備，讓我們來完成剩下的代碼吧，也就是timeout decorator：

class Timeout(Exception):
    """function run timeout"""
    
def timeout(seconds):
    """逾時裝飾器，指定逾時時間
    若被裝飾的方法在指定的時間內未返回，則拋出Timeout異常"""
    def timeout_decorator(func):
        """真正的裝飾器"""
        
        def _new_func(oldfunc, result, oldfunc_args, oldfunc_kwargs):
            result.append(oldfunc(*oldfunc_args, **oldfunc_kwargs))
        
        def _(*args, **kwargs):
            result = []
            new_kwargs = { # create new args for _new_func, because we want to get the func return val to result list
                'oldfunc': func,
                'result': result,
                'oldfunc_args': args,
                'oldfunc_kwargs': kwargs
            }
            thd = KThread(target=_new_func, args=(), kwargs=new_kwargs)
            thd.start()
            thd.join(seconds)
            alive = thd.isAlive()
            thd.kill() # kill the child thread
            if alive:
                raise Timeout(u'function run too long, timeout %d seconds.' % seconds)
            else:
                return result[0]
        _.__name__ = func.__name__
        _.__doc__ = func.__doc__
        return _
    return timeout_decorator

 

真的能運行嗎？寫個測試程式運行運行：

@timeout(5)
def method_timeout(seconds, text):
    print 'start', seconds, text
    time.sleep(seconds)
    print 'finish', seconds, text
    return seconds

if __name__ == '__main__':
    for sec in range(1, 10):
        try:
            print '*' * 20
            print method_timeout(sec, 'test waiting %d seconds' % sec)
        except Timeout, e:
            print e

 

看，真的行：

 

 

原始代碼: threadutil.py.zip

 

本次的tips可能有點複雜，運用到多線程，裝飾器等，希望它對你有所協助。

PS：山頂風光果然一流，獨缺美女相伴!!!