Skip to content

Commit 027a833

Browse files
gh-139373: Fix asyncio Process.communicate() losing output when cancelled (#154223)
Output already read when communicate() is cancelled (e.g. by a wait_for() timeout) is now retained on the Process and returned by a subsequent communicate() call, matching subprocess.Popen.communicate() behaviour on timeout. Passing input to a communicate() call following a cancelled one now raises ValueError.
1 parent daf09e1 commit 027a833

4 files changed

Lines changed: 167 additions & 5 deletions

File tree

Doc/library/asyncio-subprocess.rst

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -240,10 +240,34 @@ their completion.
240240
Note, that the data read is buffered in memory, so do not use
241241
this method if the data size is large or unlimited.
242242

243+
If this coroutine is cancelled (for example, when a timeout is
244+
set with :func:`~asyncio.wait_for`), the output that was already
245+
read is not lost: call :meth:`!communicate` again to read the
246+
remaining output and get the complete data::
247+
248+
try:
249+
stdout, stderr = await asyncio.wait_for(
250+
proc.communicate(), timeout=5.0)
251+
except TimeoutError:
252+
proc.kill()
253+
stdout, stderr = await proc.communicate()
254+
255+
Passing *input* after a previous :meth:`!communicate` call was
256+
cancelled raises :exc:`ValueError`; pass ``input=None`` to
257+
continue the communication, the original *input* is used.
258+
243259
.. versionchanged:: 3.12
244260

245261
*stdin* gets closed when ``input=None`` too.
246262

263+
.. versionchanged:: next
264+
265+
If :meth:`!communicate` is cancelled, the output that was
266+
already read is now preserved and returned by a subsequent
267+
:meth:`!communicate` call. Passing *input* to a
268+
:meth:`!communicate` call following a cancelled one now raises
269+
:exc:`ValueError`.
270+
247271
.. method:: send_signal(signal)
248272

249273
Sends the signal *signal* to the child process.

Lib/asyncio/subprocess.py

Lines changed: 30 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -124,6 +124,11 @@ def __init__(self, transport, protocol, loop):
124124
self.stdout = protocol.stdout
125125
self.stderr = protocol.stderr
126126
self.pid = transport.get_pid()
127+
self._communication_started = False
128+
self._input = None
129+
self._input_written = False
130+
self._stdout_buf = bytearray()
131+
self._stderr_buf = bytearray()
127132

128133
def __repr__(self):
129134
return f'<{self.__class__.__name__} {self.pid}>'
@@ -148,8 +153,9 @@ def kill(self):
148153
async def _feed_stdin(self, input):
149154
debug = self._loop.get_debug()
150155
try:
151-
if input is not None:
156+
if input is not None and not self._input_written:
152157
self.stdin.write(input)
158+
self._input_written = True
153159
if debug:
154160
logger.debug(
155161
'%r communicate: feed stdin (%s bytes)', self, len(input))
@@ -172,22 +178,33 @@ async def _read_stream(self, fd):
172178
transport = self._transport.get_pipe_transport(fd)
173179
if fd == 2:
174180
stream = self.stderr
181+
buf = self._stderr_buf
175182
else:
176183
assert fd == 1
177184
stream = self.stdout
185+
buf = self._stdout_buf
178186
if self._loop.get_debug():
179187
name = 'stdout' if fd == 1 else 'stderr'
180188
logger.debug('%r communicate: read %s', self, name)
181-
output = await stream.read()
189+
# Append each block to the persistent buffer as soon as it is
190+
# read so that no output is lost if this coroutine is cancelled.
191+
while block := await stream.read(stream._limit):
192+
buf += block
182193
if self._loop.get_debug():
183194
name = 'stdout' if fd == 1 else 'stderr'
184195
logger.debug('%r communicate: close %s', self, name)
185196
transport.close()
186-
return output
187197

188198
async def communicate(self, input=None):
199+
if self._communication_started:
200+
if input:
201+
raise ValueError(
202+
"Cannot send input after starting communication")
203+
else:
204+
self._input = input
205+
self._communication_started = True
189206
if self.stdin is not None:
190-
stdin = self._feed_stdin(input)
207+
stdin = self._feed_stdin(self._input)
191208
else:
192209
stdin = self._noop()
193210
if self.stdout is not None:
@@ -198,8 +215,16 @@ async def communicate(self, input=None):
198215
stderr = self._read_stream(2)
199216
else:
200217
stderr = self._noop()
201-
stdin, stdout, stderr = await tasks.gather(stdin, stdout, stderr)
218+
await tasks.gather(stdin, stdout, stderr)
202219
await self.wait()
220+
if self.stdout is not None:
221+
stdout = self._stdout_buf.take_bytes()
222+
else:
223+
stdout = None
224+
if self.stderr is not None:
225+
stderr = self._stderr_buf.take_bytes()
226+
else:
227+
stderr = None
203228
return (stdout, stderr)
204229

205230

Lib/test/test_asyncio/test_subprocess.py

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -177,6 +177,116 @@ async def run():
177177
self.assertEqual(exitcode, 0)
178178
self.assertEqual(stdout, b'')
179179

180+
def test_communicate_cancelled_mid_read_retry(self):
181+
# gh-139373: output read before communicate() was cancelled must
182+
# be returned by a subsequent communicate() call.
183+
code = ('import sys, time;'
184+
'sys.stdout.buffer.write(b"first\\n");'
185+
'sys.stdout.buffer.flush();'
186+
'time.sleep(3600)')
187+
188+
async def run():
189+
proc = await asyncio.create_subprocess_exec(
190+
sys.executable, '-c', code,
191+
stdout=subprocess.PIPE,
192+
)
193+
task = asyncio.create_task(proc.communicate())
194+
# wait until communicate() has read the first line
195+
while not proc._stdout_buf:
196+
await asyncio.sleep(0.01)
197+
task.cancel()
198+
with self.assertRaises(asyncio.CancelledError):
199+
await task
200+
proc.kill()
201+
stdout, stderr = await proc.communicate()
202+
return stdout
203+
204+
task = asyncio.wait_for(run(), support.LONG_TIMEOUT)
205+
stdout = self.loop.run_until_complete(task)
206+
self.assertEqual(stdout, b'first\n')
207+
208+
def test_communicate_cancelled_during_wait_retry(self):
209+
# gh-139373: cancellation landing after the output was fully read
210+
# but before wait() completed must not lose the output.
211+
code = ('import os, time;'
212+
'os.write(1, b"all output\\n");'
213+
'os.close(1);'
214+
'time.sleep(3600)')
215+
216+
async def run():
217+
proc = await asyncio.create_subprocess_exec(
218+
sys.executable, '-c', code,
219+
stdout=subprocess.PIPE,
220+
)
221+
task = asyncio.create_task(proc.communicate())
222+
# wait until the stdout reader has consumed everything up to
223+
# EOF; communicate() is then blocked on wait()
224+
while not proc.stdout.at_eof():
225+
await asyncio.sleep(0.01)
226+
task.cancel()
227+
with self.assertRaises(asyncio.CancelledError):
228+
await task
229+
proc.kill()
230+
stdout, stderr = await proc.communicate()
231+
return stdout
232+
233+
task = asyncio.wait_for(run(), support.LONG_TIMEOUT)
234+
stdout = self.loop.run_until_complete(task)
235+
self.assertEqual(stdout, b'all output\n')
236+
237+
def test_communicate_cancelled_input_not_resendable(self):
238+
# gh-139373: like subprocess.Popen.communicate(), sending new
239+
# input after a cancelled communicate() call is an error.
240+
async def run():
241+
proc = await asyncio.create_subprocess_exec(
242+
*PROGRAM_CAT,
243+
stdin=subprocess.PIPE,
244+
stdout=subprocess.PIPE,
245+
)
246+
task = asyncio.create_task(proc.communicate(b'data'))
247+
await asyncio.sleep(0)
248+
task.cancel()
249+
with self.assertRaises(asyncio.CancelledError):
250+
await task
251+
with self.assertRaises(ValueError):
252+
await proc.communicate(b'more data')
253+
proc.kill()
254+
await proc.communicate()
255+
256+
self.loop.run_until_complete(
257+
asyncio.wait_for(run(), support.LONG_TIMEOUT))
258+
259+
def test_communicate_cancelled_stdin_retry(self):
260+
# gh-139373: input already fed before cancellation is not re-sent
261+
# by the retried communicate() call, and the output is preserved.
262+
code = ('import sys, time;'
263+
'sys.stdout.buffer.write(sys.stdin.buffer.read());'
264+
'sys.stdout.buffer.flush();'
265+
'time.sleep(3600)')
266+
267+
async def run():
268+
proc = await asyncio.create_subprocess_exec(
269+
sys.executable, '-c', code,
270+
stdin=subprocess.PIPE,
271+
stdout=subprocess.PIPE,
272+
)
273+
task = asyncio.create_task(proc.communicate(b'hello'))
274+
# the child echoes stdin only after it is closed, so once
275+
# output arrives the input was fully written and
276+
# communicate() is blocked on wait()
277+
while not proc._stdout_buf:
278+
await asyncio.sleep(0.01)
279+
task.cancel()
280+
with self.assertRaises(asyncio.CancelledError):
281+
await task
282+
proc.kill()
283+
stdout, stderr = await proc.communicate()
284+
return stdout
285+
286+
task = asyncio.wait_for(run(), support.LONG_TIMEOUT)
287+
stdout = self.loop.run_until_complete(task)
288+
self.assertEqual(stdout, b'hello')
289+
180290
def test_shell(self):
181291
proc = self.loop.run_until_complete(
182292
asyncio.create_subprocess_shell('exit 7')
Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
Fix :meth:`asyncio.subprocess.Process.communicate` losing already-read
2+
output when it is cancelled; the output is now retained and returned by a
3+
subsequent ``communicate()`` call. Patch by Kumar Aditya.

0 commit comments

Comments
 (0)