How can I ensure closing all connection in loadbalancer when something fails or hangs?

I'm trying to write a simple load-balancer. It works ok till one of servers (BalanceServer) doesn't close connection then... Client (ReverseProxy) disconnects but the connection in with BalanceServer stays open. I tried to add callback (#3) to ReverseProxy.connectionLost to close the connection with one of the servers as I do with closing connection when server disconnects (clientLoseConnection), but at that time the ServerWriter is Null and I cannot terminate it at #1 and #2

How can I ensure that all connections are closed when one of sides disconnects? I guess that also some kind of timeout here would be nice when both client and one of servers hang, but how can I add it so it works on both connections?

from twisted.internet.protocol import Protocol, Factory, ClientCreator
from twisted.internet import reactor, defer
from collections import namedtuple

BalanceServer = namedtuple('BalanceServer', 'host port')

SERVER_LIST = [BalanceServer('127.0.0.1', 8000), BalanceServer('127.0.0.1', 8001)]

def getServer(servers):
    while True:
        for server in servers:
            yield server

# this writes to one of balance servers and responds to client with callback 'clientWrite'
class ServerWriter(Protocol):
    def sendData(self, data):
        self.transport.write(data)

    def dataReceived(self, data):
        self.clientWrite(data)

    def connectionLost( self, reason ):
        self.clientLoseConnection()

# callback for reading data from client to send it to server and get response to client again    
def transferData(serverWriter, clientWrite, clientLoseConnection, data):
    if serverWriter:
        serverWriter.clientWrite = clientWrite
        serverWriter.clientLoseConnection = clientLoseConn开发者_StackOverflow社区ection
        serverWriter.sendData(data)

def closeConnection(serverWriter):
    if serverWriter: #1 this is null
        #2 So connection is not closed and hangs there, till BalanceServer close it 
        serverWriter.transport.loseConnection()

# accepts clients
class ReverseProxy(Protocol):
    def connectionMade(self):
        server = self.factory.getServer()
        self.serverWriter = ClientCreator(reactor, ServerWriter)
        self.client = self.serverWriter.connectTCP( server.host, server.port )

    def dataReceived(self, data):
        self.client.addCallback(transferData, self.transport.write, 
                    self.transport.loseConnection, data )

    def connectionLost(self, reason):
        self.client.addCallback(closeConnection) #3 adding close doesn't work


class ReverseProxyFactory(Factory):
    protocol = ReverseProxy
    def __init__(self, serverGenerator):
        self.getServer = serverGenerator

plainFactory = ReverseProxyFactory( getServer(SERVER_LIST).next )
reactor.listenTCP( 7777, plainFactory )
reactor.run()

You may want to look at twisted.internet.protocols.portforward for an example of hooking up two connections and then disconnecting them. Or just use txloadbalancer and don't even write your own code.

However, loseConnection will never forcibly terminate the connection if there is never any traffic going over it. So if you don't have an application-level ping or any data going over your connections, they may still never shut down. This is a long-standing bug in Twisted. Actually, the longest-standing bug. Perhaps you'd like to help work on the fix :).