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

Make sure updates are not lost in sender

Browse Source
This commit is contained in:
Lonami Exo 2024-10-19 23:08:13 +02:00
parent edf3e06bd0
commit e4e7681051
2 changed files with 21 additions and 22 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
client/src/telethon/_impl/client/client/net.py
View File

@ -265,7 +265,7 @@ async def invoke_request(
async def step_sender(client: Client) -> None: async def step_sender(client: Client) -> None:
try: try:
assert client._sender assert client._sender
updates = await client._sender.step() await client._sender.step()
except ConnectionError: except ConnectionError:
if client.connected: if client.connected:
raise raise
@ -273,6 +273,7 @@ async def step_sender(client: Client) -> None:
# disconnect was called, so the socket returning 0 bytes is expected # disconnect was called, so the socket returning 0 bytes is expected
return return
updates = client._sender.pop_updates()
process_socket_updates(client, updates) process_socket_updates(client, updates)

40
client/src/telethon/_impl/mtsender/sender.py
View File

@ -169,6 +169,7 @@ class Sender:
_transport: Transport _transport: Transport
_mtp: Mtp _mtp: Mtp
_mtp_buffer: bytearray _mtp_buffer: bytearray
_updates: list[Updates]
_requests: list[Request[object]] _requests: list[Request[object]]
_request_event: Event _request_event: Event
_next_ping: float _next_ping: float
@ -200,6 +201,7 @@ class Sender:
_transport=transport, _transport=transport,
_mtp=mtp, _mtp=mtp,
_mtp_buffer=bytearray(), _mtp_buffer=bytearray(),
_updates=[],
_requests=[], _requests=[],
_request_event=Event(), _request_event=Event(),
_next_ping=asyncio.get_running_loop().time() + PING_DELAY, _next_ping=asyncio.get_running_loop().time() + PING_DELAY,
@ -236,19 +238,22 @@ class Sender:
if rx.done(): if rx.done():
return rx.result() return rx.result()
async def step(self) -> list[Updates]: async def step(self) -> None:
ticket_number = self._step_counter ticket_number = self._step_counter
async with self.lock: async with self.lock:
if self._step_counter == ticket_number: if self._step_counter == ticket_number:
# We're the one to drive IO. # We're the one to drive IO.
self._step_counter += 1 self._step_counter += 1
return await self._step() await self._step()
else: # else: different task drove IO.
# A different task drove IO.
return []
async def _step(self) -> list[Updates]: def pop_updates(self) -> list[Updates]:
updates = self._updates[:]
self._updates.clear()
return updates
async def _step(self) -> None:
self._try_fill_write() self._try_fill_write()
recv_req = asyncio.create_task(self._request_event.wait()) recv_req = asyncio.create_task(self._request_event.wait())
@ -265,16 +270,14 @@ class Sender:
task.cancel() task.cancel()
await asyncio.wait(pending) await asyncio.wait(pending)
result = []
if recv_req in done: if recv_req in done:
self._request_event.clear() self._request_event.clear()
if recv_data in done: if recv_data in done:
result = self._on_net_read(recv_data.result()) self._on_net_read(recv_data.result())
if send_data in done: if send_data in done:
self._on_net_write() self._on_net_write()
if not done: if not done:
self._on_ping_timeout() self._on_ping_timeout()
return result
async def _do_send(self) -> None: async def _do_send(self) -> None:
if self._write_drain_pending: if self._write_drain_pending:
@ -305,13 +308,12 @@ class Sender:
self._transport.pack(mtp_buffer, self._writer.write) self._transport.pack(mtp_buffer, self._writer.write)
self._write_drain_pending = True self._write_drain_pending = True
def _on_net_read(self, read_buffer: bytes) -> list[Updates]: def _on_net_read(self, read_buffer: bytes) -> None:
if not read_buffer: if not read_buffer:
raise ConnectionResetError("read 0 bytes") raise ConnectionResetError("read 0 bytes")
self._read_buffer += read_buffer self._read_buffer += read_buffer
updates: list[Updates] = []
while self._read_buffer: while self._read_buffer:
self._mtp_buffer.clear() self._mtp_buffer.clear()
try: try:
@ -320,9 +322,7 @@ class Sender:
break break
else: else:
del self._read_buffer[:n] del self._read_buffer[:n]
self._process_mtp_buffer(updates) self._process_mtp_buffer()
return updates
def _on_net_write(self) -> None: def _on_net_write(self) -> None:
for req in self._requests: for req in self._requests:
@ -340,12 +340,12 @@ class Sender:
) )
self._next_ping = asyncio.get_running_loop().time() + PING_DELAY self._next_ping = asyncio.get_running_loop().time() + PING_DELAY
def _process_mtp_buffer(self, updates: list[Updates]) -> None: def _process_mtp_buffer(self) -> None:
results = self._mtp.deserialize(self._mtp_buffer) results = self._mtp.deserialize(self._mtp_buffer)
for result in results: for result in results:
if isinstance(result, Update): if isinstance(result, Update):
self._process_update(updates, result.body) self._process_update(result.body)
elif isinstance(result, RpcResult): elif isinstance(result, RpcResult):
self._process_result(result) self._process_result(result)
elif isinstance(result, RpcError): elif isinstance(result, RpcError):
@ -353,11 +353,9 @@ class Sender:
else: else:
self._process_bad_message(result) self._process_bad_message(result)
def _process_update( def _process_update(self, update: bytes | bytearray | memoryview) -> None:
self, updates: list[Updates], update: bytes | bytearray | memoryview
) -> None:
try: try:
updates.append(Updates.from_bytes(update)) self._updates.append(Updates.from_bytes(update))
except ValueError: except ValueError:
cid = struct.unpack_from("I", update)[0] cid = struct.unpack_from("I", update)[0]
alt_classes: tuple[Type[Serializable], ...] = ( alt_classes: tuple[Type[Serializable], ...] = (
@ -377,7 +375,7 @@ class Sender:
AffectedMessages, AffectedMessages,
), ),
) )
updates.append( self._updates.append(
UpdateShort( UpdateShort(
update=UpdateDeleteMessages( update=UpdateDeleteMessages(
messages=[], messages=[],