From 592e107f52165295bd531f9e01756ce22b88c722 Mon Sep 17 00:00:00 2001 From: Lonami Exo Date: Tue, 29 Aug 2017 13:35:56 +0200 Subject: [PATCH] Recreate the TCP socket only when needed --- telethon/extensions/tcp_client.py | 26 +++++++++++++++----------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/telethon/extensions/tcp_client.py b/telethon/extensions/tcp_client.py index 5abf620d..2e0a1e25 100644 --- a/telethon/extensions/tcp_client.py +++ b/telethon/extensions/tcp_client.py @@ -11,9 +11,8 @@ from ..errors import ReadCancelledError class TcpClient: def __init__(self, proxy=None): - self.connected = False self._proxy = proxy - self._recreate_socket() + self._socket = None # Support for multi-threading advantages and safety self.cancelled = Event() # Has the read operation been cancelled? @@ -36,23 +35,27 @@ class TcpClient: 'timeout' must be given in seconds """ if not self.connected: + self._recreate_socket() self._socket.settimeout(timeout) self._socket.connect((ip, port)) self._socket.setblocking(False) - self.connected = True + + def _get_connected(self): + return self._socket is not None + + connected = property(fget=_get_connected) def close(self): """Closes the connection""" - if self.connected: - try: + try: + if self.connected: self._socket.shutdown(socket.SHUT_RDWR) self._socket.close() - except OSError as e: - if e.errno != errno.ENOTCONN: - raise - - self.connected = False - self._recreate_socket() + except OSError as e: + if e.errno != errno.ENOTCONN: + raise + finally: + self._socket = None def write(self, data): """Writes (sends) the specified bytes to the connected peer""" @@ -66,6 +69,7 @@ class TcpClient: try: sent = self._socket.send(view[total_sent:]) if sent == 0: + self.close() raise ConnectionResetError( 'The server has closed the connection.') total_sent += sent