LonamiWebs/Telethon
1
1
Fork 0
mirror of https://github.com/LonamiWebs/Telethon.git synced 2025-03-03 19:00:21 +03:00
Code Issues Packages Projects Releases Wiki Activity

Recreate the TCP socket only when needed

Browse Source
This commit is contained in:
Lonami Exo 2017-08-29 13:35:56 +02:00
parent cff7655a10
commit 592e107f52
1 changed files with 15 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

26
telethon/extensions/tcp_client.py
View File

@ -11,9 +11,8 @@ from ..errors import ReadCancelledError
class TcpClient: class TcpClient:
def __init__(self, proxy=None): def __init__(self, proxy=None):
self.connected = False
self._proxy = proxy self._proxy = proxy
self._recreate_socket() self._socket = None
# Support for multi-threading advantages and safety # Support for multi-threading advantages and safety
self.cancelled = Event() # Has the read operation been cancelled? self.cancelled = Event() # Has the read operation been cancelled?
@ -36,23 +35,27 @@ class TcpClient:
'timeout' must be given in seconds 'timeout' must be given in seconds
""" """
if not self.connected: if not self.connected:
self._recreate_socket()
self._socket.settimeout(timeout) self._socket.settimeout(timeout)
self._socket.connect((ip, port)) self._socket.connect((ip, port))
self._socket.setblocking(False) 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): def close(self):
"""Closes the connection""" """Closes the connection"""
if self.connected: try:
try: if self.connected:
self._socket.shutdown(socket.SHUT_RDWR) self._socket.shutdown(socket.SHUT_RDWR)
self._socket.close() self._socket.close()
except OSError as e: except OSError as e:
if e.errno != errno.ENOTCONN: if e.errno != errno.ENOTCONN:
raise raise
finally:
self.connected = False self._socket = None
self._recreate_socket()
def write(self, data): def write(self, data):
"""Writes (sends) the specified bytes to the connected peer""" """Writes (sends) the specified bytes to the connected peer"""
@ -66,6 +69,7 @@ class TcpClient:
try: try:
sent = self._socket.send(view[total_sent:]) sent = self._socket.send(view[total_sent:])
if sent == 0: if sent == 0:
self.close()
raise ConnectionResetError( raise ConnectionResetError(
'The server has closed the connection.') 'The server has closed the connection.')
total_sent += sent total_sent += sent