cache.py

from twisted.internet import reactor
from twisted.protocols.memcache import MemCacheProtocol
from txconnpool.pool import PooledClientFactory, Pool


class _PooledMemCacheProtocol(MemCacheProtocol):
    factory = None

    def connectionMade(self):
        MemCacheProtocol.connectionMade(self)

        self.factory.connectionPool.clientFree(self)

        if self.factory.deferred is not None:
            self.factory.deferred.callback(self)
            self.factory.deferred = None


class _MemCacheClientFactory(PooledClientFactory):
    protocol = _PooledMemCacheProtocol


class MemCachePool(Pool):
    clientFactory = _MemCacheClientFactory

    def get(self, *args, **kwargs):
        return self.perform_request_with_timeout('get', *args, **kwargs)

    def set(self, *args, **kwargs):
        return self.perform_request_with_timeout('set', *args, **kwargs)

    def perform_request_with_timeout(self, request, *args, **kwargs):
        def _timeout_conn_check():
            d.cancel()

        dc = reactor.callLater(2, _timeout_conn_check)

        def _cancel_timeout(result):
            if dc.active():
                dc.cancel()

            return result

        d = self.performRequest(request, *args, **kwargs)
        d.addBoth(_cancel_timeout)
        return d