Skip to content

Commit

Permalink
[Fabric-Sync] Do not use stdout with non-blocking asyncio (#35446)
Browse files Browse the repository at this point in the history
Wrapping stdout with StreamWriter leads to setting stdout stream to
non-blocking. In consequence that parent process and all child processes
have stdio in non-blocking mode, which is not a standard setup. This
leads to random failures (e.g. Python's print might trow BlockinIOError
exception).
  • Loading branch information
arkq authored and pull[bot] committed Nov 15, 2024
1 parent 3bfe87b commit 2069810
Showing 1 changed file with 23 additions and 40 deletions.
63 changes: 23 additions & 40 deletions examples/fabric-admin/scripts/fabric-sync-app.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,30 +20,13 @@
import shutil
import signal
import sys
import typing
from argparse import ArgumentParser
from tempfile import TemporaryDirectory


async def asyncio_stdin() -> asyncio.StreamReader:
"""Wrap sys.stdin in an asyncio StreamReader."""
loop = asyncio.get_event_loop()
reader = asyncio.StreamReader()
protocol = asyncio.StreamReaderProtocol(reader)
await loop.connect_read_pipe(lambda: protocol, sys.stdin)
return reader


async def asyncio_stdout(file=sys.stdout) -> asyncio.StreamWriter:
"""Wrap an IO stream in an asyncio StreamWriter."""
loop = asyncio.get_event_loop()
transport, protocol = await loop.connect_write_pipe(
lambda: asyncio.streams.FlowControlMixin(loop=loop),
os.fdopen(file.fileno(), 'wb'))
return asyncio.streams.StreamWriter(transport, protocol, None, loop)


async def forward_f(prefix: bytes, f_in: asyncio.StreamReader,
f_out: asyncio.StreamWriter, cb=None):
f_out: typing.BinaryIO, cb=None):
"""Forward f_in to f_out with a prefix attached.
This function can optionally feed received lines to a callback function.
Expand All @@ -54,9 +37,9 @@ async def forward_f(prefix: bytes, f_in: asyncio.StreamReader,
break
if cb is not None:
cb(line)
f_out.write(prefix)
f_out.write(line)
await f_out.drain()
f_out.buffer.write(prefix)
f_out.buffer.write(line)
f_out.flush()


async def forward_pipe(pipe_path: str, f_out: asyncio.StreamWriter):
Expand All @@ -72,6 +55,7 @@ async def forward_pipe(pipe_path: str, f_out: asyncio.StreamWriter):
data = os.read(fd, 1024)
if data:
f_out.write(data)
await f_out.drain()
if not data:
await asyncio.sleep(0.1)
except BlockingIOError:
Expand All @@ -80,13 +64,17 @@ async def forward_pipe(pipe_path: str, f_out: asyncio.StreamWriter):

async def forward_stdin(f_out: asyncio.StreamWriter):
"""Forward stdin to f_out."""
reader = await asyncio_stdin()
loop = asyncio.get_event_loop()
reader = asyncio.StreamReader()
protocol = asyncio.StreamReaderProtocol(reader)
await loop.connect_read_pipe(lambda: protocol, sys.stdin)
while True:
line = await reader.readline()
if not line:
# Exit on Ctrl-D (EOF).
sys.exit(0)
f_out.write(line)
await f_out.drain()


class Subprocess:
Expand All @@ -109,15 +97,9 @@ async def run(self):
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE)
# Add the stdout and stderr processing to the event loop.
asyncio.create_task(forward_f(
self.tag,
self.p.stderr,
await asyncio_stdout(sys.stderr)))
asyncio.create_task(forward_f(
self.tag,
self.p.stdout,
await asyncio_stdout(sys.stdout),
cb=self._check_output))
asyncio.create_task(forward_f(self.tag, self.p.stderr, sys.stderr))
asyncio.create_task(forward_f(self.tag, self.p.stdout, sys.stdout,
cb=self._check_output))

async def send(self, message: str, expected_output: str = None, timeout: float = None):
"""Send a message to a process and optionally wait for a response."""
Expand Down Expand Up @@ -206,14 +188,6 @@ async def main(args):
if pipe and not os.path.exists(pipe):
os.mkfifo(pipe)

def terminate(signum, frame):
admin.terminate()
bridge.terminate()
sys.exit(0)

signal.signal(signal.SIGINT, terminate)
signal.signal(signal.SIGTERM, terminate)

admin, bridge = await asyncio.gather(
run_admin(
args.app_admin,
Expand All @@ -235,6 +209,15 @@ def terminate(signum, frame):
passcode=args.passcode,
))

def terminate():
admin.terminate()
bridge.terminate()
sys.exit(0)

loop = asyncio.get_event_loop()
loop.add_signal_handler(signal.SIGINT, terminate)
loop.add_signal_handler(signal.SIGTERM, terminate)

# Wait a bit for apps to start.
await asyncio.sleep(1)

Expand Down

0 comments on commit 2069810

Please sign in to comment.