pool.py

class Pool:
    def __init__(self,
                 connection_factory,
                 target_size=0,
                 decay_time=0):
        self.connection_factory = connection_factory
        self.target_size = target_size
        self.decay_time = decay_time
        self._connections = {}
        self.pool_lock = RLock()
        self.last_cleanup = monotonic()

        for _ in range(self.target_size):
            with self.pool_lock:
                self._make_connection()

    def _get_idle_connection(self):
        connection_keys = list(self._connections.keys())
        shuffle(connection_keys)
        for key in connection_keys:
            conn = self._connections[key]
            refcount = weakref.getweakrefcount(conn)
            if refcount == 0:
                return key, weakref.proxy(conn)
        return None, None

    def _get_connection(self):
        key, conn = self._get_idle_connection()
        if conn is not None:
            return key, conn
        return self._make_connection()

    def _abandon_extra_connection(self):
        if monotonic() - self.last_cleanup > self.decay_time:
            if len(self._connections) > self.target_size:
                key, conn = self._get_idle_connection()
                if key is not None:
                    log.info('Removing extra db connection %s',
                             key)
                    del self._connections[key]
            else:
                self.last_cleanup = monotonic()

    def get_connection(self):
        for attempt in range(100):
            log.debug('Trying to get a working db connection, try %s',
                      attempt)
            with self.pool_lock:
                key, conn = self._get_connection()
            try:
                conn.execute('SELECT 1 AS foo;')
                conn.rollback()
            except Exception:
                log.exception('Connection %s: %s has died',
                              key,
                              conn)
                with self.pool_lock:
                    del self._connections[key]
            break

        with self.pool_lock:
            self._abandon_extra_connection()
        return conn

    def _make_connection(self):
        log.info('Creating DB connection')
        key = uuid4()
        conn = self.connection_factory()
        with self.pool_lock:
            self._connections[key] = conn

        return key, weakref.proxy(conn)