_GeneratorTask: throw CancelledError in cancelled coroutine (bug 711174)

Throw asyncio.CancelledError in a cancelled coroutine, ensuring
that the coroutine can handle this exception in order to perform
any necessary cleanup (like close the log file for bug 711174).
Note that the asyncio.CancelledError will only be thrown in the
coroutine if there's an opportunity (yield) before the generator
raises StopIteration.

Also fix the AsynchronousTask exit listener handling for
compatibility with this new behavior.

Fixes: 8074127bbc21 ("SpawnProcess: add _main coroutine")
Bug: https://bugs.gentoo.org/711174
Signed-off-by: Zac Medico <zmedico@gentoo.org>

3 files changed, 49 insertions, 11 deletions

diff --git a/lib/_emerge/AsynchronousTask.py b/lib/_emerge/AsynchronousTask.py
index 1e9e177cb..580eef050 100644
--- a/lib/_emerge/AsynchronousTask.py
+++ b/lib/_emerge/AsynchronousTask.py
@@ -64,7 +64,7 @@ class AsynchronousTask(SlotObject):
 		@returns: Future, result is self.returncode
 		"""
 		waiter = self.scheduler.create_future()
-		exit_listener = lambda self: waiter.set_result(self.returncode)
+		exit_listener = lambda self: waiter.cancelled() or waiter.set_result(self.returncode)
 		self.addExitListener(exit_listener)
 		waiter.add_done_callback(lambda waiter:
 			self.removeExitListener(exit_listener) if waiter.cancelled() else None)
@@ -180,9 +180,15 @@ class AsynchronousTask(SlotObject):
 	def removeExitListener(self, f):
 		if self._exit_listeners is None:
 			if self._exit_listener_stack is not None:
-				self._exit_listener_stack.remove(f)
+				try:
+					self._exit_listener_stack.remove(f)
+				except ValueError:
+					pass
 			return
-		self._exit_listeners.remove(f)
+		try:
+			self._exit_listeners.remove(f)
+		except ValueError:
+			pass
 
 	def _wait_hook(self):
 		"""
diff --git a/lib/portage/tests/util/futures/test_compat_coroutine.py b/lib/portage/tests/util/futures/test_compat_coroutine.py
index f96aa9be5..b561c0227 100644
--- a/lib/portage/tests/util/futures/test_compat_coroutine.py
+++ b/lib/portage/tests/util/futures/test_compat_coroutine.py
@@ -57,20 +57,43 @@ class CompatCoroutineTestCase(TestCase):
 			loop.run_until_complete(catching_coroutine(loop=loop)))
 
 	def test_cancelled_coroutine(self):
+		"""
+		Verify that a coroutine can handle (and reraise) asyncio.CancelledError
+		in order to perform any necessary cleanup. Note that the
+		asyncio.CancelledError will only be thrown in the coroutine if there's
+		an opportunity (yield) before the generator raises StopIteration.
+		"""
+		loop = asyncio.get_event_loop()
+		ready_for_exception = loop.create_future()
+		exception_in_coroutine = loop.create_future()
 
 		@coroutine
 		def cancelled_coroutine(loop=None):
 			loop = asyncio._wrap_loop(loop)
 			while True:
-				yield loop.create_future()
+				task = loop.create_future()
+				try:
+					ready_for_exception.set_result(None)
+					yield task
+				except BaseException as e:
+					# Since python3.8, asyncio.CancelledError inherits
+					# from BaseException.
+					task.done() or task.cancel()
+					exception_in_coroutine.set_exception(e)
+					raise
+				else:
+					exception_in_coroutine.set_result(None)
 
-		loop = asyncio.get_event_loop()
 		future = cancelled_coroutine(loop=loop)
-		loop.call_soon(future.cancel)
+		loop.run_until_complete(ready_for_exception)
+		future.cancel()
 
 		self.assertRaises(asyncio.CancelledError,
 			loop.run_until_complete, future)
 
+		self.assertRaises(asyncio.CancelledError,
+			loop.run_until_complete, exception_in_coroutine)
+
 	def test_cancelled_future(self):
 		"""
 		When a coroutine raises CancelledError, the coroutine's
diff --git a/lib/portage/util/futures/compat_coroutine.py b/lib/portage/util/futures/compat_coroutine.py
index b745fd845..54fc316fe 100644
--- a/lib/portage/util/futures/compat_coroutine.py
+++ b/lib/portage/util/futures/compat_coroutine.py
@@ -87,21 +87,29 @@ class _GeneratorTask(object):
 	def __init__(self, generator, result, loop):
 		self._generator = generator
 		self._result = result
+		self._current_task = None
 		self._loop = loop
 		result.add_done_callback(self._cancel_callback)
 		loop.call_soon(self._next)
 
 	def _cancel_callback(self, result):
-		if result.cancelled():
-			self._generator.close()
+		if result.cancelled() and self._current_task is not None:
+			# The done callback for self._current_task invokes
+			# _next in either case here.
+			self._current_task.done() or self._current_task.cancel()
 
 	def _next(self, previous=None):
+		self._current_task = None
 		if self._result.cancelled():
 			if previous is not None:
 				# Consume exceptions, in order to avoid triggering
 				# the event loop's exception handler.
 				previous.cancelled() or previous.exception()
-			return
+
+			# This will throw asyncio.CancelledError in the coroutine if
+			# there's an opportunity (yield) before the generator raises
+			# StopIteration.
+			previous = self._result
 		try:
 			if previous is None:
 				future = next(self._generator)
@@ -124,5 +132,6 @@ class _GeneratorTask(object):
 			if not self._result.cancelled():
 				self._result.set_exception(e)
 		else:
-			future = asyncio.ensure_future(future, loop=self._loop)
-			future.add_done_callback(self._next)
+			self._current_task = asyncio.ensure_future(future, loop=self._loop)
+			self._current_task.add_done_callback(self._next)
+