coroutine with custom timeout implementation - python-3.x

Question

I'm using asyncio to run a subprocess that prints some output which I parse line by line and do various things based on the output I see. I want to put a timeout on this process, but it should not be a global timeout for the life of the entire process. Instead, whenever I see certain specific output from the process, I want to actually reset the timeout so that it starts over. How can I implement this?
For a global timeout I have this working and it's easy, I simply call asyncio.wait_for(_foo(), timeout). But I can't get this to work with resetting the timeout. Here's what I have so far:
# here invocation is my own data structure with some bookkeeping information in it
# (such as the start time from which I want to base my timeout decisions on).
# and process is the value returned by await asyncio.create_subprocess_exec(...)
# _run_one_invocation() is my own function which is just a stdout readline loop
# and some dispatching.
# Make a Task out of the co-routine so that we decide when it gets cancelled, not Python.
run_task = asyncio.Task(_run_one_invocation(invocation, process))
while True:
try:
# Use asyncio.shield so it doesn't get cancelled if the timeout expires
await asyncio.shield(asyncio.wait_for(run_task, 1))
# If the await returns without raising an exception, we got to EOF and we're done.
break
except asyncio.TimeoutError:
# If it's been too long since last reset, this is a "real" timeout.
duration = time.time() - invocation.start_time
if duration > timeout:
run_task.cancel()
raise
When I run this, the if statement which calls run_task.cancel() is not being entered, and nevertheless when I go back to the top of the loop and call asyncio.wait_for() again it immediately raises an asyncio.CancelledError.
What am I doing wrong?

Answer 1

You can fix the issue and simplify the code by completely avoiding wait_for() (and therefore shield()) and just using wait(return_when=FIRST_COMPLETED) to implement the kind of timeout you need:
run_task = asyncio.create_task(_run_one_invocation(invocation, process))
while True:
await asyncio.wait([run_task], timeout=1)
if run_task.done():
break
if time.time() - invocation.start_time > timeout:
run_task.cancel()
raise asyncio.TimeoutErrror()
The downside of this approach is that it introduces 1s wakeups, prohibiting the program (and consequently the computer) from ever going to sleep, even if the task is dormant for hours. Probably not a big deal on a server, but such practices contribute to battery drain on laptops, and it's a good idea to avoid them. Also, the 1s sleep introduces an up to 1s latency to react to a change in timeout.
To resolve that, you can create an event that is fired by the code changing the timeout, and react to that event in addition to the timeout and the task completing:
timeout_changed = asyncio.Event()
# pass timeout_changed where needed, and have the code that changes
# the timeout also call timeout_changed.set()
run_task = asyncio.create_task(_run_one_invocation(invocation, process))
while True:
remaining = timeout - (time.time() - invocation.start_time)
timeout_changed_task = asyncio.ensure_future(timeout_changed.wait())
await asyncio.wait([run_task, timeout_changed_task],
return_when=asyncio.FIRST_COMPLETED, timeout=remaining)
timeout_changed_task.cancel()
# either: 1) the task has completed, 2) the previous timeout has
# expired, or 3) the timeout has changed
if run_task.done():
break # 1
if time.time() - invocation.start_time > timeout:
# 2 or 2+3
run_task.cancel()
raise asyncio.TimeoutErrror()
# 3 - continue waiting with the new timeout