boto / s3transfer

Amazon S3 Transfer Manager for Python
Apache License 2.0
209 stars 133 forks source link

Fix hang on shutdown with multiple threadpools #266

Open TedSinger opened 1 year ago

TedSinger commented 1 year ago

This patch fixes an issue that occurs when a download completion is expected after the threadpools have shutdown.

Here is a test that reproduces the problem. Without the patch, running this hangs forever, demonstrating the bug. With the patch, pytest terminates as soon as the KeyboardInterrupt is thrown, as expected. This test is not appropriate to include in the test suite, since it halts the whole test runner on success and hangs on failure. (Any error lesser than KeyboardInterrupt fails to demonstrate the bug under pytest. pytest catches the error and prevents the global shutdown hook from triggering until the tests are cleaned up)

    # in tests.integration.test_download.TestDownload
    def test_concurrent_shutdown_in_exception(self):
        import io
        from concurrent.futures import ThreadPoolExecutor

        names = ['foo.txt', 'bar.txt']
        for name in names:
            filename = self.files.create_file_with_size(
                name, filesize=750 * 1024 * 1024
            )
            self.upload_file(filename name)

        def download(key):
            # This is how a boto3 client implements `download_fileobj`
            with self.create_transfer_manager(self.config) as manager:
                future = manager.download(
                    bucket=self.bucket_name,
                    key=key
                    fileobj=io.BytesIO()
                )
                return future.result()

        mapper = ThreadPoolExecutor(max_workers=2)
        iterable = mapper.map(download, names)
        try:
            # Suppose some unrelated error causes Python to try to shutdown...
            raise KeyboardInterrupt()
        finally:
            # Because we are running things concurrently, some other threadpool
            # demands a new result, and we wait on an already started future. That
            # should be fine, as the future should either finish or crash.
            # But s3transfer.futures.TransferCoordinator fails to call self.announce_done()
            # and hangs forever
            next(iterable)

This patch is licensed Apache 2.0, in accordance with the boto/s3transfer contribution policy.