Skis / Toys / Fun

If you have a few minutes and speak python & twisted, it would be useful to have an extra set of eyes on this section of code. The basic idea of this is to be a reconnecting thrift client, such that I can just write simple client.function(a,b,c) calls without having to worry about if there is or isn’t a client and it will queue reconnect as needed.

from thrift.transport import TTwisted
from thrift.protocol import TBinaryProtocol
from twisted.internet.protocol import ReconnectingClientFactory
from twisted.internet import defer, reactor
from twisted.python import failure
from collections import deque
from redback import log

class ClientBusy(Exception):
    pass

class ClientDead(Exception):
    pass

class InvalidThriftRequest(Exception):
    pass

class ManagedThriftRequest(object):
    def __init__(self, method, *args, **kw) :
        self.method = method
        self.args   = args
        self.kw     = kw

class ManagedClient(object) :
    def __init__(self, factory) :
        self.__factory = factory

    def _is_connected(self) :
        return self.__factory.is_connected()

    def __getattr__(self, name) :
        if hasattr(self.__factory.client_class, name) :
            def f(*args, **kw) :
                return self.__factory.pushRequest(ManagedThriftRequest(name, *args, **kw))
            return f
        raise InvalidThriftRequest("Cant find method: %s" % name)

class ManagedThriftClientProtocol(TTwisted.ThriftClientProtocol):
    def __init__(self, client_class, iprot_factory, oprot_factory=None):
        TTwisted.ThriftClientProtocol.__init__(self, client_class, iprot_factory, oprot_factory)
        self.client_class = client_class
        self.deferred = None
        self.alive    = False

    def connectionMade(self):
        log.debug(self, "Connection made to", "%s.%s" % (self.client_class.__module__, self.client_class.__name__), "[%s]:%s" % self.transport.addr)
        self.alive = True
        TTwisted.ThriftClientProtocol.connectionMade(self)
        self.client.protocol = self
        self.factory.clientIdle(self)

    def connectionLost(self, reason=None):
        log.debug(self, "Connection lost to", "%s.%s" % (self.client_class.__module__, self.client_class.__name__), "[%s]:%s" % self.transport.addr)
        self.alive = False
        try :
            TTwisted.ThriftClientProtocol.connectionLost(self, reason)
            self.factory.clientGone(self)
        except Exception, e :
            log.error(self.connectionLost, e)

    def _complete(self, res, request, dfd):
        self.deferred = None
        if isinstance(res, failure.Failure) and dfd :
            self.factory.pushRequest(request, dfd)
        else :
            if dfd :
                dfd.callback(res)
            self.factory.clientIdle(self)
        return res

    def submitRequest(self, request, dfd):
        if not self.alive :
            raise ClientBusy
        if not self.deferred :
            fun = getattr(self.client, request.method, None)
            if not fun:
                raise InvalidThriftRequest("No such method as : %s" % request.method)
            else :
                try :
                    d = fun(*request.args, **request.kw)
                except Exception, e :
                    log.error(self.submitRequest, "calling : ", request.method, e)
            self.deferred = d
            d.addBoth(self._complete, request, dfd)
            return d
        else:
            raise ClientBusy

class ManagedClientFactory(ReconnectingClientFactory):
    maxDelay        = 5
    thriftFactory   = TBinaryProtocol.TBinaryProtocolAcceleratedFactory
    protocol        = ManagedThriftClientProtocol
    submitLoopSleep = 0
    client_class    = None

    def __init__(self, client_class=None):
        self._stack        = deque()
        self._protos       = defer.DeferredQueue()
        self.deferred      = defer.Deferred()
        self.client_class  = client_class or self.client_class
        self.client        = ManagedClient(self)

    def _errback(self, reason=None):
        if self.deferred :
            self.deferred.errback(reason)
            self.deferred = None

    def _callback(self, value=None):
        if self.deferred :
            self.deferred.callback(value)
            self.deferred = None

    def clientConnectionFailed(self, connector, reason):
        ReconnectingClientFactory.clientConnectionFailed(self, connector, reason)
        self._errback(reason)

    def buildProtocol(self, addr) :
        self.resetDelay()
        p = self.protocol(self.client_class, self.thriftFactory())

        p.factory = self
        #self._protos.put(p)
        return p

    def clientIdle(self, proto) :
        self._callback(True)
        if proto.alive :
            self._protos.put(proto)

    def clientGone(self, proto):
        pass
        #if proto in self._protos :
        #    self._protos.remove(proto)

    def _protoErr(self, proto):
        pass
        #import traceback
        #traceback.print_stack()
        #print "IN PROTO ERR", proto

    def _protoReady(self, proto):
        if proto.deferred :
            log.msg(self._protoReady, "Proto currently active!")
            return
        if not proto.alive :
            log.msg(self._protoReady, "Proto currently dead!")
            return

        try:
            request, deferred = self._stack.popleft()
        except defer.QueueUnderflow :
            pass

        d = proto.submitRequest(request, deferred)
        return d

    def pushRequest(self, request, din=None) :
        d = din or defer.Deferred()
        self._stack.append((request, d))

        dfd = self._protos.get()
        dfd.addCallback(self._protoReady)
        dfd.addErrback(self._protoErr)
        return d

    def shutdown(self) :
        """Shutdown this factory"""

        self.stopTrying()
        for p in self._protos:
            if p.transport:
                p.transport.loseConnection()</pre>