bpo-30064: Fix asyncio loop.sock_* race condition issue (#20369)

This commit is contained in:
Fantix King 2020-05-27 14:47:30 -05:00 committed by GitHub
parent 526e23f153
commit 210a137396
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 157 additions and 16 deletions

View File

@ -266,6 +266,7 @@ class BaseSelectorEventLoop(base_events.BaseEventLoop):
(handle, writer))
if reader is not None:
reader.cancel()
return handle
def _remove_reader(self, fd):
if self.is_closed():
@ -302,6 +303,7 @@ class BaseSelectorEventLoop(base_events.BaseEventLoop):
(reader, handle))
if writer is not None:
writer.cancel()
return handle
def _remove_writer(self, fd):
"""Remove a writer callback."""
@ -329,7 +331,7 @@ class BaseSelectorEventLoop(base_events.BaseEventLoop):
def add_reader(self, fd, callback, *args):
"""Add a reader callback."""
self._ensure_fd_no_transport(fd)
return self._add_reader(fd, callback, *args)
self._add_reader(fd, callback, *args)
def remove_reader(self, fd):
"""Remove a reader callback."""
@ -339,7 +341,7 @@ class BaseSelectorEventLoop(base_events.BaseEventLoop):
def add_writer(self, fd, callback, *args):
"""Add a writer callback.."""
self._ensure_fd_no_transport(fd)
return self._add_writer(fd, callback, *args)
self._add_writer(fd, callback, *args)
def remove_writer(self, fd):
"""Remove a writer callback."""
@ -362,13 +364,15 @@ class BaseSelectorEventLoop(base_events.BaseEventLoop):
pass
fut = self.create_future()
fd = sock.fileno()
self.add_reader(fd, self._sock_recv, fut, sock, n)
self._ensure_fd_no_transport(fd)
handle = self._add_reader(fd, self._sock_recv, fut, sock, n)
fut.add_done_callback(
functools.partial(self._sock_read_done, fd))
functools.partial(self._sock_read_done, fd, handle=handle))
return await fut
def _sock_read_done(self, fd, fut):
self.remove_reader(fd)
def _sock_read_done(self, fd, fut, handle=None):
if handle is None or not handle.cancelled():
self.remove_reader(fd)
def _sock_recv(self, fut, sock, n):
# _sock_recv() can add itself as an I/O callback if the operation can't
@ -401,9 +405,10 @@ class BaseSelectorEventLoop(base_events.BaseEventLoop):
pass
fut = self.create_future()
fd = sock.fileno()
self.add_reader(fd, self._sock_recv_into, fut, sock, buf)
self._ensure_fd_no_transport(fd)
handle = self._add_reader(fd, self._sock_recv_into, fut, sock, buf)
fut.add_done_callback(
functools.partial(self._sock_read_done, fd))
functools.partial(self._sock_read_done, fd, handle=handle))
return await fut
def _sock_recv_into(self, fut, sock, buf):
@ -446,11 +451,12 @@ class BaseSelectorEventLoop(base_events.BaseEventLoop):
fut = self.create_future()
fd = sock.fileno()
fut.add_done_callback(
functools.partial(self._sock_write_done, fd))
self._ensure_fd_no_transport(fd)
# use a trick with a list in closure to store a mutable state
self.add_writer(fd, self._sock_sendall, fut, sock,
memoryview(data), [n])
handle = self._add_writer(fd, self._sock_sendall, fut, sock,
memoryview(data), [n])
fut.add_done_callback(
functools.partial(self._sock_write_done, fd, handle=handle))
return await fut
def _sock_sendall(self, fut, sock, view, pos):
@ -502,9 +508,11 @@ class BaseSelectorEventLoop(base_events.BaseEventLoop):
# connection runs in background. We have to wait until the socket
# becomes writable to be notified when the connection succeed or
# fails.
self._ensure_fd_no_transport(fd)
handle = self._add_writer(
fd, self._sock_connect_cb, fut, sock, address)
fut.add_done_callback(
functools.partial(self._sock_write_done, fd))
self.add_writer(fd, self._sock_connect_cb, fut, sock, address)
functools.partial(self._sock_write_done, fd, handle=handle))
except (SystemExit, KeyboardInterrupt):
raise
except BaseException as exc:
@ -512,8 +520,9 @@ class BaseSelectorEventLoop(base_events.BaseEventLoop):
else:
fut.set_result(None)
def _sock_write_done(self, fd, fut):
self.remove_writer(fd)
def _sock_write_done(self, fd, fut, handle=None):
if handle is None or not handle.cancelled():
self.remove_writer(fd)
def _sock_connect_cb(self, fut, sock, address):
if fut.done():

View File

@ -1,4 +1,5 @@
import socket
import time
import asyncio
import sys
from asyncio import proactor_events
@ -122,6 +123,136 @@ class BaseSockTestsMixin:
sock = socket.socket()
self._basetest_sock_recv_into(httpd, sock)
async def _basetest_sock_recv_racing(self, httpd, sock):
sock.setblocking(False)
await self.loop.sock_connect(sock, httpd.address)
task = asyncio.create_task(self.loop.sock_recv(sock, 1024))
await asyncio.sleep(0)
task.cancel()
asyncio.create_task(
self.loop.sock_sendall(sock, b'GET / HTTP/1.0\r\n\r\n'))
data = await self.loop.sock_recv(sock, 1024)
# consume data
await self.loop.sock_recv(sock, 1024)
self.assertTrue(data.startswith(b'HTTP/1.0 200 OK'))
async def _basetest_sock_recv_into_racing(self, httpd, sock):
sock.setblocking(False)
await self.loop.sock_connect(sock, httpd.address)
data = bytearray(1024)
with memoryview(data) as buf:
task = asyncio.create_task(
self.loop.sock_recv_into(sock, buf[:1024]))
await asyncio.sleep(0)
task.cancel()
task = asyncio.create_task(
self.loop.sock_sendall(sock, b'GET / HTTP/1.0\r\n\r\n'))
nbytes = await self.loop.sock_recv_into(sock, buf[:1024])
# consume data
await self.loop.sock_recv_into(sock, buf[nbytes:])
self.assertTrue(data.startswith(b'HTTP/1.0 200 OK'))
await task
async def _basetest_sock_send_racing(self, listener, sock):
listener.bind(('127.0.0.1', 0))
listener.listen(1)
# make connection
sock.setblocking(False)
task = asyncio.create_task(
self.loop.sock_connect(sock, listener.getsockname()))
await asyncio.sleep(0)
server = listener.accept()[0]
server.setblocking(False)
with server:
await task
# fill the buffer
with self.assertRaises(BlockingIOError):
while True:
sock.send(b' ' * 5)
# cancel a blocked sock_sendall
task = asyncio.create_task(
self.loop.sock_sendall(sock, b'hello'))
await asyncio.sleep(0)
task.cancel()
# clear the buffer
async def recv_until():
data = b''
while not data:
data = await self.loop.sock_recv(server, 1024)
data = data.strip()
return data
task = asyncio.create_task(recv_until())
# immediately register another sock_sendall
await self.loop.sock_sendall(sock, b'world')
data = await task
# ProactorEventLoop could deliver hello
self.assertTrue(data.endswith(b'world'))
async def _basetest_sock_connect_racing(self, listener, sock):
listener.bind(('127.0.0.1', 0))
addr = listener.getsockname()
sock.setblocking(False)
task = asyncio.create_task(self.loop.sock_connect(sock, addr))
await asyncio.sleep(0)
task.cancel()
listener.listen(1)
i = 0
while True:
try:
await self.loop.sock_connect(sock, addr)
break
except ConnectionRefusedError: # on Linux we need another retry
await self.loop.sock_connect(sock, addr)
break
except OSError as e: # on Windows we need more retries
# A connect request was made on an already connected socket
if getattr(e, 'winerror', 0) == 10056:
break
# https://stackoverflow.com/a/54437602/3316267
if getattr(e, 'winerror', 0) != 10022:
raise
i += 1
if i >= 128:
raise # too many retries
# avoid touching event loop to maintain race condition
time.sleep(0.01)
def test_sock_client_racing(self):
with test_utils.run_test_server() as httpd:
sock = socket.socket()
with sock:
self.loop.run_until_complete(asyncio.wait_for(
self._basetest_sock_recv_racing(httpd, sock), 10))
sock = socket.socket()
with sock:
self.loop.run_until_complete(asyncio.wait_for(
self._basetest_sock_recv_into_racing(httpd, sock), 10))
listener = socket.socket()
sock = socket.socket()
with listener, sock:
self.loop.run_until_complete(asyncio.wait_for(
self._basetest_sock_send_racing(listener, sock), 10))
listener = socket.socket()
sock = socket.socket()
with listener, sock:
self.loop.run_until_complete(asyncio.wait_for(
self._basetest_sock_connect_racing(listener, sock), 10))
async def _basetest_huge_content(self, address):
sock = socket.socket()
sock.setblocking(False)

View File

@ -0,0 +1 @@
Fix asyncio ``loop.sock_*`` race condition issue