Added a thread lock to the TcpClient

This gives multi-threading safety without giving up on speed
(now there's no need for additional sleeps)
This commit is contained in:
Lonami 2016-09-10 10:17:15 +02:00
parent fd14a5a49a
commit e47344c0f0
3 changed files with 51 additions and 36 deletions

View File

@ -14,22 +14,27 @@ from tl.all_tlobjects import tlobjects
class MtProtoSender: class MtProtoSender:
"""MTProto Mobile Protocol sender (https://core.telegram.org/mtproto/description)""" """MTProto Mobile Protocol sender (https://core.telegram.org/mtproto/description)"""
def __init__(self, transport, session): def __init__(self, transport, session, check_updates_delay=0.1):
"""If check_updates_delay is None, no updates will be checked.
Otherwise, specifies every how often updates should be checked"""
self.transport = transport self.transport = transport
self.session = session self.session = session
self.need_confirmation = [] # Message IDs that need confirmation self.need_confirmation = [] # Message IDs that need confirmation
self.on_update_handlers = [] self.on_update_handlers = []
# Set up updates thread # Set up updates thread, if the delay is not None
self.updates_thread = Thread(target=self.updates_thread_method, name='Updates thread') self.check_updates_delay = check_updates_delay
self.updates_thread_running = True if check_updates_delay:
self.updates_thread_paused = True self.updates_thread = Thread(target=self.updates_thread_method, name='Updates thread')
self.updates_thread_running = True
self.updates_thread_paused = True
self.updates_thread.start() self.updates_thread.start()
def disconnect(self): def disconnect(self):
"""Disconnects and **stops all the running threads**""" """Disconnects and **stops all the running threads** if any"""
self.updates_thread_running = False self.updates_thread_running = False
self.transport.cancel_receive() self.transport.cancel_receive()
self.transport.close() self.transport.close()

View File

@ -1,6 +1,7 @@
# Python rough implementation of a C# TCP client # Python rough implementation of a C# TCP client
import socket import socket
import time import time
from threading import Lock
from errors import ReadCancelledError from errors import ReadCancelledError
from utils import BinaryWriter from utils import BinaryWriter
@ -10,8 +11,11 @@ class TcpClient:
def __init__(self): def __init__(self):
self.connected = False self.connected = False
self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
# Support for multi-threading advantages and safety
self.cancelled = False # Has the read operation been cancelled? self.cancelled = False # Has the read operation been cancelled?
self.delay = 0.1 # Read delay when there was no data available self.delay = 0.1 # Read delay when there was no data available
self.lock = Lock()
def connect(self, ip, port): def connect(self, ip, port):
"""Connects to the specified IP and port number""" """Connects to the specified IP and port number"""
@ -27,40 +31,46 @@ class TcpClient:
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"""
self.socket.sendall(data)
# Ensure that only one thread can send data at once
with self.lock:
self.socket.sendall(data)
def read(self, buffer_size): def read(self, buffer_size):
"""Reads (receives) the specified bytes from the connected peer""" """Reads (receives) the specified bytes from the connected peer"""
self.cancelled = False # Ensure it is not cancelled at first
with BinaryWriter() as writer: # Ensure that only one thread can receive data at once
while writer.written_count < buffer_size and not self.cancelled: with self.lock:
try: # Ensure it is not cancelled at first, so we can enter the loop
# When receiving from the socket, we may not receive all the data at once self.cancelled = False
# This is why we need to keep checking to make sure that we receive it all
left_count = buffer_size - writer.written_count
partial = self.socket.recv(left_count)
writer.write(partial)
except BlockingIOError: with BinaryWriter() as writer:
# There was no data available for us to read. Sleep a bit while writer.written_count < buffer_size and not self.cancelled:
time.sleep(self.delay) try:
# When receiving from the socket, we may not receive all the data at once
# This is why we need to keep checking to make sure that we receive it all
left_count = buffer_size - writer.written_count
partial = self.socket.recv(left_count)
writer.write(partial)
# If the operation was cancelled *but* data was read, except BlockingIOError:
# this will result on data loss so raise an exception # There was no data available for us to read. Sleep a bit
# TODO this could be solved by using an internal FIFO buffer (first in, first out) time.sleep(self.delay)
if self.cancelled:
if writer.written_count == 0:
raise ReadCancelledError()
else:
raise NotImplementedError('The read operation was cancelled when some data '
'was already read. This has not yet implemented '
'an internal buffer, so cannot continue.')
return writer.get_bytes() # If the operation was cancelled *but* data was read,
# this will result on data loss so raise an exception
# TODO this could be solved by using an internal FIFO buffer (first in, first out)
if self.cancelled:
if writer.written_count == 0:
raise ReadCancelledError()
else:
raise NotImplementedError('The read operation was cancelled when some data '
'was already read. This has not yet implemented '
'an internal buffer, so cannot continue.')
# If everything went fine, return the read bytes
return writer.get_bytes()
def cancel_read(self): def cancel_read(self):
"""Cancels the read operation if it was blocking and stops """Cancels the read operation raising a ReadCancelledError"""
the current thread until it's cancelled"""
self.cancelled = True self.cancelled = True
time.sleep(self.delay)

View File

@ -60,8 +60,8 @@ class TcpTransport:
self.tcp_client.close() self.tcp_client.close()
def cancel_receive(self): def cancel_receive(self):
"""Cancels (stops) trying to receive from the remote peer and """Cancels (stops) trying to receive from the
stops the current thread until it's cancelled""" remote peer and raises a ReadCancelledError"""
self.tcp_client.cancel_read() self.tcp_client.cancel_read()
def get_client_delay(self): def get_client_delay(self):