Commit 59430898 authored by Antoine Millet's avatar Antoine Millet
Browse files

Added close_endpoint_on_shutdown option on TunnelProtocol

Also removed endpoint close from default on shutdown callback.
parent e19773ee
Loading
Loading
Loading
Loading
+7 −3
Original line number Diff line number Diff line
@@ -18,7 +18,8 @@ class TunnelProtocol(Protocol):
    def __init__(self, *args, **kwargs):
        endpoint = kwargs.pop('endpoint', None)
        autostart = kwargs.pop('autostart', True)
        self._cb_on_shutdown = kwargs.pop('on_shutdown', self.cb_default_on_shutdown)
        self._cb_on_shutdown = kwargs.pop('on_shutdown', None)
        self._close_endpoint_on_shutdown = kwargs.pop('close_endpoint_on_shutdown', True)
        super(TunnelProtocol, self).__init__(*args, **kwargs)

        if endpoint is None:
@@ -112,7 +113,7 @@ class TunnelProtocol(Protocol):
    def cb_default_on_shutdown(self, tun):
        """ Action to do on the endpoint when the connection is shutdown.
        """
        tun.endpoint.close()
        pass

#
# Public methods:
@@ -151,7 +152,10 @@ class TunnelProtocol(Protocol):
            except RpcConnectionError:
                pass  # Ignore errors when connection is down.
            # Execute the callback:
            if self._cb_on_shutdown is not None:
                self._cb_on_shutdown(self)
            if self._close_endpoint_on_shutdown:
                self.endpoint.close()
            super(TunnelProtocol, self).shutdown()

    def wait_shutdown(self, timeout=None, sig_timeout=1):