Ticket #1237: pool.py

# -*- coding: gb2312 -*-
# 
# author: junzhang.jn@gmail.com

import threading,thread
import time
import sys,traceback

class DBPoolWithThread:
    # @param factory      PooledObjectFactory object create factory
    # @param freetime     integer             when object's idle time > freetime , check_thread will destory it
    # @param threadsafety integer             like dbapi2.0��
    #                                               0 = Threads may not share the module. 
    #                                               1 = Threads may share the module, but not connections. 
    #                                               2 = Threads may share the module and connections. 
    #                                               3 = Threads may share the module, connections and cursors. 
    def __init__( self , factory , freetime = 30 , threadsafety = 1 ):
        self.factory = factory
        self.threadsafety = threadsafety
        if threadsafety == 0:
            raise RuntimeError( 'threadsafety is %d , you cannot use connectionpool for this db driver' % threadsafety )
        self.idlepool = {}   # { threadid: [] , }
        self.lock = threading.Lock()
        self.freetime = freetime
        self.idleTime = {} # { obj: ( threadid , time ) , }
        self.deleted = False
        self.exitEvent = threading.Event()
        self.thread = threading.Thread( target=DBPoolWithThread.check_thread , args = ( self , ) )
        self.thread.start()
    
    def __container_get_obj( self ):
        try:
            self.lock.acquire()
            threadid = thread.get_ident()
            obj = None
            try:
                # 1.1 get object in current thread 
                li = None
                try:
                    li = self.idlepool[ threadid ]
                    # 1.2 get first object
                    obj = li.pop(0)
                    if len( li ) == 0:
                        del self.idlepool[ threadid ]
                except KeyError , e:
                    # not find object list in current thread
                    if self.threadsafety == 1:
                        return None
                except IndexError , e :
                    # this exception should not occured.
                    del self.idlepool[ threadid ]
                
                # 1.3 not found object in current thread
                if obj == None and self.threadsafety > 1:
                    # 2 get object from other thread, threadsafety MUST > 1
                    for key , cons in self.idlepool.items():
                        obj = cons.pop(0)
                        if len( cons ) == 0:
                            del self.idlepool[ key ]
                        break
            except :
                pass
            if obj and self.idleTime.has_key( obj ):
                del self.idleTime[ obj ]    # clear idle timer for return object
            return obj
        finally:
            self.lock.release()
    
    def __container_put_obj( self , obj ):
        try:
            self.lock.acquire()
            threadid = thread.get_ident()
            try:
                li = self.idlepool[ threadid ]
            except KeyError , e:
                li = []
                self.idlepool[ threadid ] =  li
            li.append( obj )
            self.idleTime[ obj ] = ( threadid , time.time() )
        finally:
            self.lock.release()
    
    def borrow_object( self ):
        i = 0
        while i < 5:
            obj = self.__container_get_obj()
            
            if not obj :
                obj = self.factory.create_object()
                break
            
            if self.factory.validate_object( obj ):
                # object is valid
                break
            else:
                # not valid
                self.factory.destroy_object( obj )# destroy invalid object
                obj = None
            i += 1
        return obj
    
    def return_object( self , obj ):
        if obj:
            self.__container_put_obj( obj )
    
    def __del__( self ):
        try:
            self.lock.acquire()
            self.exitEvent.set()
            for key , cons in self.idlepool.items():
                for obj in cons:
                    self.factory.destroy_object( obj )
        finally:
            self.lock.release()
    
    def check_thread( this ):
        while True:
            this.exitEvent.wait( this.freetime )    # wait interval time
            if this.exitEvent.isSet():
                break
            if this.freetime:
                this.lock.acquire()
                try:
                    threadid = thread.get_ident()
                    for k,v in this.idleTime.items():
                        if time.time() - v[1] > this.freetime: # timeout will destroy
                            li = this.idlepool[ v[0] ]
                            li.remove( k )
                            this.factory.destroy_object( k )
                            if len( li ) == 0:
                                del this.idlepool[ v[0] ] # delete list
                            del this.idleTime[ k ]
                finally:
                    this.lock.release()
    check_thread = staticmethod( check_thread )
    
# object create factory interface
class PooledObjectFactory:
    def __init__( self ):
        pass
        
    def create_object( self ):
        raise NotImplementedError()
    
    def destroy_object( self , obj ):
        raise NotImplementedError()
    
    def validate_object( self , obj ):
        raise NotImplementedError()