Cancel futures if concurrent.futures.Executor is used as a context manager and there is an exception?

I bumped into a bit of a rough edge related to how concurrent.futures.Executor instances behave when they're used as context managers and there is an exception raised from the with statement body.

I had some code like this (reduced to the bare minimum):

with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
    futures = {
        executor.submit(check_status, item): item
        for item in to_check
    }
    for index, future in enumerate(concurrent.futures.as_completed(futures)):
        print(f"{index}/{len(futures)} {future.result()}")

Now, as it happens check_status liked to raise an exception from time to time and this would be fine if the application stopped because of that but it kept going but in a weird way. What I saw was check_status kept producing logs but the print calls in the snippet above stopped being executed.

This is because of how Executor implements __exit__ – it waits for all pending futures to complete before propagating the exception:

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.shutdown(wait=True)
        return False

No in my case I started with thousands of futures so it could take hours for the process to finally crash.

Since Python 3.9 the shutdown method has the cancel_futures parameter (implemented in https://github.com/python/cpython/pull/18057): https://docs.python.org/3.9/library/concurrent.futures.html#concurrent.futures.Executor.shutdown.

I'm proposing changing the behavior of the context manager __exit__ to cancel futures if there is an exception:

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.shutdown(wait=True, cancel_futures=exc_val is not None)
        return False

This would allow the in-flight futures to finish but any further processing wouldn't happen and the exception would be reported in a (more) timely fashion.

cc @brianquinlan @pitrou as I saw you participate in the cancel_futures PR, if there's someone else I should ping here for visibility please let me know.