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

Prefer BufferedWriter over BinaryWriter for socket operations

Browse Source
This commit is contained in:
Lonami Exo 2017-06-09 11:42:39 +02:00
parent 5494cb15eb
commit c6645a555d
2 changed files with 25 additions and 21 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

41
telethon/network/tcp_client.py
View File

@ -2,10 +2,10 @@
import socket import socket
import time import time
from datetime import datetime, timedelta from datetime import datetime, timedelta
from io import BytesIO, BufferedWriter
from threading import Event, Lock from threading import Event, Lock
from ..errors import ReadCancelledError from ..errors import ReadCancelledError
from ..utils import BinaryWriter
class TcpClient: class TcpClient:
@ -58,40 +58,44 @@ class TcpClient:
'The server has closed the connection.') 'The server has closed the connection.')
total_sent += sent total_sent += sent
def read(self, buffer_size, timeout=timedelta(seconds=5)): def read(self, size, timeout=timedelta(seconds=5)):
"""Reads (receives) the specified bytes from the connected peer. """Reads (receives) a whole block of 'size bytes
A timeout can be specified, which will cancel the operation if no data from the connected peer.
has been read in the specified time. If data was read and it's waiting
for more, the timeout will NOT cancel the operation. Set to None for no timeout""" A timeout can be specified, which will cancel the operation if
no data has been read in the specified time. If data was read
and it's waiting for more, the timeout will NOT cancel the
operation. Set to None for no timeout
"""
# Ensure that only one thread can receive data at once # Ensure that only one thread can receive data at once
with self._lock: with self._lock:
# Ensure it is not cancelled at first, so we can enter the loop # Ensure it is not cancelled at first, so we can enter the loop
self.cancelled.clear() self.cancelled.clear()
# Set the starting time so we can calculate whether the timeout should fire # Set the starting time so we can
if timeout: # calculate whether the timeout should fire
start_time = datetime.now() start_time = datetime.now() if timeout else None
with BinaryWriter() as buffer: with BufferedWriter(BytesIO(), buffer_size=size) as buffer:
while buffer.written_count < buffer_size: bytes_left = size
while bytes_left != 0:
# Only do cancel if no data was read yet # Only do cancel if no data was read yet
# Otherwise, carry on reading and finish # Otherwise, carry on reading and finish
if self.cancelled.is_set() and buffer.written_count == 0: if self.cancelled.is_set() and bytes_left == size:
raise ReadCancelledError() raise ReadCancelledError()
try: try:
# When receiving from the socket, we may not receive all the data at once partial = self._socket.recv(bytes_left)
# This is why we need to keep checking to make sure that we receive it all
left_count = buffer_size - buffer.written_count
partial = self._socket.recv(left_count)
if len(partial) == 0: if len(partial) == 0:
raise ConnectionResetError( raise ConnectionResetError(
'The server has closed the connection (recv() returned 0 bytes).') 'The server has closed the connection (recv() returned 0 bytes).')
buffer.write(partial) buffer.write(partial)
bytes_left -= len(partial)
except BlockingIOError as error: except BlockingIOError as error:
# There was no data available for us to read. Sleep a bit # No data available yet, sleep a bit
time.sleep(self.delay) time.sleep(self.delay)
# Check if the timeout finished # Check if the timeout finished
@ -102,7 +106,8 @@ class TcpClient:
'The read operation exceeded the timeout.') from error 'The read operation exceeded the timeout.') from error
# If everything went fine, return the read bytes # If everything went fine, return the read bytes
return buffer.get_bytes() buffer.flush()
return buffer.raw.getvalue()
def cancel_read(self): def cancel_read(self):
"""Cancels the read operation IF it hasn't yet """Cancels the read operation IF it hasn't yet

5
telethon/utils/binary_writer.py
View File

@ -12,8 +12,7 @@ class BinaryWriter:
if not stream: if not stream:
stream = BytesIO() stream = BytesIO()
self.stream = stream self.writer = BufferedWriter(stream)
self.writer = BufferedWriter(self.stream)
self.written_count = 0 self.written_count = 0
# region Writing # region Writing
@ -126,7 +125,7 @@ class BinaryWriter:
"""Get the current bytes array content from the buffer, optionally flushing first""" """Get the current bytes array content from the buffer, optionally flushing first"""
if flush: if flush:
self.writer.flush() self.writer.flush()
return self.stream.getvalue() return self.writer.raw.getvalue()
def get_written_bytes_count(self): def get_written_bytes_count(self):
"""Gets the count of bytes written in the buffer. """Gets the count of bytes written in the buffer.