TeamPGM/pyrogram
2
0
Fork 0
mirror of https://github.com/TeamPGM/pyrogram.git synced 2024-11-27 16:45:19 +00:00
Code Issues Packages Projects Releases Wiki Activity

Always use non-blocking sockets

Browse Source
This commit is contained in:
Dan 2022-12-26 16:25:22 +01:00
parent 9ee1807e42
commit 6aae3a9c77
1 changed files with 11 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
pyrogram/connection/transport/tcp/tcp.py
View File

@ -42,8 +42,11 @@ class TCP:
def __init__(self, ipv6: bool, proxy: dict): def __init__(self, ipv6: bool, proxy: dict):
self.socket = None self.socket = None
self.reader = None # type: asyncio.StreamReader self.reader = None
self.writer = None # type: asyncio.StreamWriter self.writer = None
self.send_queue = asyncio.Queue()
self.send_task = None
self.loop = asyncio.get_event_loop() self.loop = asyncio.get_event_loop()
@ -75,16 +78,16 @@ class TCP:
else socket.AF_INET else socket.AF_INET
) )
self.socket.settimeout(TCP.TIMEOUT) self.socket.setblocking(False)
self.send_queue = asyncio.Queue()
self.send_task = None
async def connect(self, address: tuple): async def connect(self, address: tuple):
await asyncio.get_event_loop().sock_connect(self.socket, address) try:
await asyncio.wait_for(asyncio.get_event_loop().sock_connect(self.socket, address), TCP.TIMEOUT)
except asyncio.TimeoutError: # Re-raise as TimeoutError. asyncio.TimeoutError is deprecated in 3.11
raise TimeoutError("Connection timed out")
self.reader, self.writer = await asyncio.open_connection(sock=self.socket) self.reader, self.writer = await asyncio.open_connection(sock=self.socket)
self.send_task = asyncio.create_task(self.send_worker()) self.send_task = asyncio.create_task(self.send_worker())
self.socket.setblocking(False)
async def close(self): async def close(self):
await self.send_queue.put(None) await self.send_queue.put(None)