From c77f45e13d80d451b06b34b618296453df66b17c Mon Sep 17 00:00:00 2001 From: Steve Yoo Date: Wed, 11 Nov 2026 16:30:38 -0500 Subject: [PATCH] Release finished CRT transfers instead of keeping them until shutdown --- .changes/next-release/bugfix-crt-72249.json | 5 + awscli/s3transfer/crt.py | 43 ++- tests/functional/s3transfer/test_crt.py | 348 ++++++++++++++++++++ tests/unit/s3transfer/test_crt.py | 11 + 4 files changed, 399 insertions(+), 8 deletions(-) create mode 100644 .changes/next-release/bugfix-crt-72249.json diff --git a/.changes/next-release/bugfix-crt-72249.json b/.changes/next-release/bugfix-crt-72249.json new file mode 100644 index 000000000000..c2b8ff67dfd3 --- /dev/null +++ b/.changes/next-release/bugfix-crt-72249.json @@ -0,0 +1,5 @@ +{ + "type": "bugfix", + "category": "crt", + "description": "Fix memory growth in ``s3`` commands that transfer many files with the CRT transfer client by releasing each transfer when it finishes instead of when the command exits." +} diff --git a/awscli/s3transfer/crt.py b/awscli/s3transfer/crt.py index f4f0b52e0740..6ba854791f33 100644 --- a/awscli/s3transfer/crt.py +++ b/awscli/s3transfer/crt.py @@ -490,7 +490,8 @@ def __init__( self._region_redirect_policy = CRTS3RegionRedirectPolicy( crt_request_serializer ) - self._future_coordinators = [] + self._future_coordinators = set() + self._future_coordinators_lock = threading.Lock() self._semaphore = threading.Semaphore(128) # not configurable # A counter to create unique id's for each transfer submitted. self._id_counter = 0 @@ -562,17 +563,29 @@ def _validate_checksum_algorithm_supported(self, extra_args): f'Supported algorithms are: {supported_algorithms}' ) + def _tracked_coordinators(self): + with self._future_coordinators_lock: + return list(self._future_coordinators) + + def _track_coordinator(self, coordinator): + with self._future_coordinators_lock: + self._future_coordinators.add(coordinator) + + def _untrack_coordinator(self, coordinator): + with self._future_coordinators_lock: + self._future_coordinators.discard(coordinator) + def _cancel_transfers(self): - for coordinator in self._future_coordinators: + for coordinator in self._tracked_coordinators(): if not coordinator.done(): coordinator.cancel() def _finish_transfers(self): - for coordinator in self._future_coordinators: + for coordinator in self._tracked_coordinators(): coordinator.result() def _wait_transfers_done(self): - for coordinator in self._future_coordinators: + for coordinator in self._tracked_coordinators(): coordinator.wait_until_on_done_callbacks_complete() def _shutdown(self, cancel=False): @@ -619,6 +632,10 @@ def _submit_transfer(self, request_type, call_args): future = CRTTransferFuture(**components) afterdone = AfterDoneHandler(coordinator) on_done_after_calls.append(afterdone) + # Untrack only after the done callbacks, which shutdown waits on. + on_done_after_calls.append( + lambda **kwargs: self._untrack_coordinator(coordinator) + ) # Serialization can rewrite an ARN in call_args, so retain the # caller-provided bucket for redirect eligibility and caching. bucket = call_args.bucket @@ -637,6 +654,7 @@ def _submit_transfer(self, request_type, call_args): is_replayable = False try: + self._track_coordinator(coordinator) self._semaphore.acquire() on_queued = self._s3_args_creator.get_crt_callback( future, 'queued' @@ -684,13 +702,18 @@ def create_request(is_region_redirect): is_replayable=is_replayable, ) except Exception as e: + # The done callbacks run on this thread, so shutdown needn't wait. + self._untrack_coordinator(coordinator) coordinator.set_exception(e, True) on_done = self._s3_args_creator.get_crt_callback( future, 'done', after_subscribers=on_done_after_calls ) coordinator.complete(e) on_done(error=e) - self._future_coordinators.append(coordinator) + except BaseException: + # An interrupted submit may never complete, so don't wait on it. + self._untrack_coordinator(coordinator) + raise self._id_counter += 1 return future @@ -1414,6 +1437,9 @@ def set_s3_request(self, s3_request, is_region_redirect=False): # The redirect is already active. The original request # completed and redirected before make_request() returned. return + if self._completion_started: + # The request finished before make_request() returned. + return if is_region_redirect: self._redirect_retry_started = True self._s3_request = s3_request @@ -1432,10 +1458,11 @@ def complete(self, error=None): self._completion_started = True completion_future = self._completion_future # No further request will be started for this transfer, so release - # the factory. It holds the transfer's call args, including an - # upload's stream, and the manager keeps every coordinator until - # it shuts down. + # the factory and the CRT request. The CRT request's native + # callbacks reference this coordinator, so keeping it would leak + # both. self._request_factory = None + self._s3_request = None if error is None: completion_future.set_result(None) else: diff --git a/tests/functional/s3transfer/test_crt.py b/tests/functional/s3transfer/test_crt.py index aa29d0ee3978..fbd88997a897 100644 --- a/tests/functional/s3transfer/test_crt.py +++ b/tests/functional/s3transfer/test_crt.py @@ -11,9 +11,11 @@ # ANY KIND, either express or implied. See the License for the specific # language governing permissions and limitations under the License. import fnmatch +import gc import io import threading import time +import weakref from concurrent.futures import Future import pytest @@ -1408,6 +1410,283 @@ def test_cancel(self): future.result() self.assertTrue(self.cancel_called) + def _assert_shutdown_returns(self, cancel=False): + shutdown = threading.Thread( + target=self.transfer_manager.shutdown, + kwargs={'cancel': cancel}, + daemon=True, + ) + shutdown.start() + shutdown.join(5) + self.assertFalse(shutdown.is_alive()) + + def _assert_coordinator_released(self, coordinator, *clients): + # The mocks keep the requests' callbacks, which reference the + # coordinator. + for client in clients or (self.s3_crt_client,): + client.reset_mock() + client.make_request.side_effect = None + # A redirected transfer finishes on its own thread. + deadline = time.time() + RESULT_TIMEOUT + while coordinator() is not None and time.time() < deadline: + gc.collect() + time.sleep(0.01) + self.assertIsNone(coordinator()) + + def _assert_result_raises(self, future): + # Not assertRaises, which keeps the traceback and so the coordinator. + raised = False + try: + future.result() + except Exception: + raised = True + self.assertTrue(raised) + + def test_finished_transfer_is_released(self): + future = self.transfer_manager.upload( + self.filename, self.bucket, self.key, {}, [] + ) + future.result() + coordinator = weakref.ref(future._coordinator) + del future + self._assert_coordinator_released(coordinator) + + def test_shutdown_waits_for_in_flight_transfer(self): + self.s3_crt_client.make_request.side_effect = None + self.s3_crt_client.make_request.return_value = self.s3_request + finished = self.transfer_manager.upload( + self.filename, self.bucket, self.key, {}, [] + ) + in_flight = self.transfer_manager.upload( + self.filename, self.bucket, self.key, {}, [] + ) + first_on_done = self.s3_crt_client.make_request.call_args_list[0][1][ + 'on_done' + ] + first_on_done(error=None) + self.assertTrue(finished.done()) + + shutdown = threading.Thread( + target=self.transfer_manager.shutdown, daemon=True + ) + shutdown.start() + shutdown.join(0.2) + self.assertTrue(shutdown.is_alive()) + + self._invoke_done_callbacks() + shutdown.join(5) + self.assertFalse(shutdown.is_alive()) + self.assertTrue(in_flight.done()) + + def test_shutdown_after_interrupted_submit_does_not_hang(self): + # Ctrl-C usually lands while a submit waits for a free request slot. + self.transfer_manager._semaphore = mock.Mock() + self.transfer_manager._semaphore.acquire.side_effect = ( + KeyboardInterrupt() + ) + with self.assertRaises(KeyboardInterrupt): + self.transfer_manager.upload( + self.filename, self.bucket, self.key, {}, [] + ) + + self._assert_shutdown_returns(cancel=True) + + def test_failed_submit_is_released(self): + self.s3_crt_client.make_request.side_effect = ( + awscrt.exceptions.from_code(0) + ) + future = self.transfer_manager.upload( + self.filename, self.bucket, self.key, {}, [] + ) + self._assert_result_raises(future) + coordinator = weakref.ref(future._coordinator) + del future + self._assert_coordinator_released(coordinator) + + def test_transfer_that_fails_in_crt_is_released(self): + self.s3_crt_client.make_request.side_effect = None + self.s3_crt_client.make_request.return_value = self.s3_request + future = self.transfer_manager.upload( + self.filename, self.bucket, self.key, {}, [] + ) + self._invoke_done_callbacks(error=awscrt.exceptions.from_code(0)) + self._assert_result_raises(future) + coordinator = weakref.ref(future._coordinator) + del future + self._assert_coordinator_released(coordinator) + + def test_redirected_transfer_is_released(self): + self.transfer_manager = self._create_redirecting_transfer_manager( + self._fail_make_request( + self._create_redirect_error('eu-central-1') + ) + ) + future = self._upload_and_wait(self.transfer_manager) + coordinator = weakref.ref(future._coordinator) + del future + self._assert_coordinator_released( + coordinator, self.initial_client, self.redirected_client + ) + + def test_shutdown_after_interrupted_seekable_check(self): + with open(self.filename, 'rb') as fileobj: + with mock.patch( + 's3transfer.crt.seekable', side_effect=KeyboardInterrupt() + ): + with self.assertRaises(KeyboardInterrupt): + self.transfer_manager.upload( + fileobj, self.bucket, self.key, {}, [] + ) + self._assert_shutdown_returns(cancel=True) + + def test_shutdown_after_failed_submit_interrupted_in_on_done(self): + # Ctrl-C can land while a failed submit runs its done callbacks. + class InterruptingSubscriber(BaseSubscriber): + def on_done(self, future, **kwargs): + raise KeyboardInterrupt() + + self.s3_crt_client.make_request.side_effect = RuntimeError() + with self.assertRaises(KeyboardInterrupt): + self.transfer_manager.upload( + self.filename, + self.bucket, + self.key, + {}, + [InterruptingSubscriber()], + ) + self._assert_shutdown_returns(cancel=True) + + def test_shutdown_after_failed_submit_with_raising_on_done(self): + class RaisingSubscriber(BaseSubscriber): + def on_done(self, future, **kwargs): + raise ValueError() + + self.s3_crt_client.make_request.side_effect = RuntimeError() + with self.assertRaises(ValueError): + self.transfer_manager.upload( + self.filename, + self.bucket, + self.key, + {}, + [RaisingSubscriber()], + ) + self._assert_shutdown_returns(cancel=True) + + def _upload_in_flight(self, count): + # As in the CRT, cancelling a request finishes it with an error. + on_dones = [] + + def make_request(**kwargs): + on_done = kwargs['on_done'] + on_dones.append(on_done) + request = mock.Mock(awscrt.s3.S3Request) + request.cancel.side_effect = lambda: on_done( + error=awscrt.exceptions.from_code(0) + ) + return request + + self.s3_crt_client.make_request.side_effect = make_request + futures = [ + self.transfer_manager.upload( + self.filename, self.bucket, self.key, {}, [] + ) + for _ in range(count) + ] + return futures, on_dones + + def test_shutdown_while_transfers_finish_their_done_callbacks(self): + futures, on_dones = self._upload_in_flight(2) + # Both transfers have a result, but their done callbacks are still + # running on CRT threads. + for future in futures: + future._coordinator.complete() + coordinator_class = s3transfer.crt.CRTTransferCoordinator + wait = coordinator_class.wait_until_on_done_callbacks_complete + + def finish_done_callbacks_then_wait(coordinator, timeout=None): + while on_dones: + on_dones.pop()(error=None) + wait(coordinator, timeout) + + with mock.patch.object( + coordinator_class, + 'wait_until_on_done_callbacks_complete', + autospec=True, + side_effect=finish_done_callbacks_then_wait, + ): + self.transfer_manager.shutdown() + + def test_interrupted_shutdown_cancels_transfers_after_others_finish(self): + futures, on_dones = self._upload_in_flight(2) + on_done_by_coordinator = { + future._coordinator: on_done + for future, on_done in zip(futures, on_dones) + } + result = s3transfer.crt.CRTTransferCoordinator.result + result_calls = [] + + def finish_first_then_interrupt(coordinator, timeout=None): + result_calls.append(coordinator) + if len(result_calls) == 1: + # The transfer shutdown waits on finishes first. + on_done_by_coordinator[coordinator](error=None) + return result(coordinator, timeout) + # Ctrl-C while shutdown waits on the next transfer. + raise KeyboardInterrupt() + + def finish_stranded_transfers(): + for coordinator, on_done in on_done_by_coordinator.items(): + if not coordinator.done(): + on_done(error=None) + + self.addCleanup(finish_stranded_transfers) + shutdown = threading.Thread( + target=self.transfer_manager.shutdown, daemon=True + ) + with mock.patch.object( + s3transfer.crt.CRTTransferCoordinator, + 'result', + autospec=True, + side_effect=finish_first_then_interrupt, + ): + shutdown.start() + shutdown.join(5) + self.assertFalse(shutdown.is_alive()) + self.assertTrue(all(future.done() for future in futures)) + + def test_shutdown_waits_for_subscriber_done_callbacks(self): + self.s3_crt_client.make_request.side_effect = None + self.s3_crt_client.make_request.return_value = self.s3_request + on_done_started = threading.Event() + release_on_done = threading.Event() + + class BlockingSubscriber(BaseSubscriber): + def on_done(self, future, **kwargs): + on_done_started.set() + release_on_done.wait(RESULT_TIMEOUT) + + self.addCleanup(release_on_done.set) + self.transfer_manager.upload( + self.filename, self.bucket, self.key, {}, [BlockingSubscriber()] + ) + on_done = threading.Thread( + target=self._invoke_done_callbacks, daemon=True + ) + on_done.start() + self.addCleanup(on_done.join, RESULT_TIMEOUT) + self.assertTrue(on_done_started.wait(RESULT_TIMEOUT)) + + shutdown = threading.Thread( + target=self.transfer_manager.shutdown, daemon=True + ) + shutdown.start() + shutdown.join(0.2) + self.assertTrue(shutdown.is_alive()) + + release_on_done.set() + shutdown.join(RESULT_TIMEOUT) + self.assertFalse(shutdown.is_alive()) + def test_serializer_error_handling(self): class SerializationException(Exception): pass @@ -1441,6 +1720,75 @@ def test_crt_s3_client_error_handling(self): future.result() +class FakeNativeS3Request: + def __init__(self, native_callbacks, key): + weakref.finalize(self, native_callbacks.pop, key, None) + + def cancel(self): + pass + + +class NativeLikeS3Client: + """Like awscrt, keeps each request's callbacks until the request is freed + and finishes requests on another thread. + """ + + def __init__(self): + self.native_callbacks = {} + self.threads = [] + self.release = threading.Event() + self._next_key = 0 + self._lock = threading.Lock() + + def make_request(self, **kwargs): + with self._lock: + key = self._next_key + self._next_key += 1 + self.native_callbacks[key] = kwargs['on_done'] + request = FakeNativeS3Request(self.native_callbacks, key) + thread = threading.Thread(target=self._finish, args=(key,)) + self.threads.append(thread) + thread.start() + return request + + def _finish(self, key): + self.release.wait(RESULT_TIMEOUT) + self.native_callbacks[key](error=None) + + +@requires_crt() +class TestCRTTransferLifetime(unittest.TestCase): + def setUp(self): + self.files = FileCreator() + self.addCleanup(self.files.remove_all) + self.filename = self.files.create_file('myfile', 'content') + session = Session() + session.set_config_variable('region', 'us-west-2') + self.client = NativeLikeS3Client() + self.transfer_manager = s3transfer.crt.CRTTransferManager( + crt_client_factory=lambda region=None: self.client, + crt_request_serializer=( + s3transfer.crt.BotocoreCRTRequestSerializer(session) + ), + ) + + def test_finished_transfers_release_coordinator_and_request(self): + coordinators = [] + for i in range(50): + future = self.transfer_manager.upload( + self.filename, 'bucket', f'key{i}', {}, [] + ) + coordinators.append(weakref.ref(future._coordinator)) + del future + self.client.release.set() + for thread in self.client.threads: + thread.join(RESULT_TIMEOUT) + gc.collect() + self.assertEqual([c for c in coordinators if c() is not None], []) + self.assertEqual(self.client.native_callbacks, {}) + self.transfer_manager.shutdown() + + MULTIPART_THRESHOLD = 8 * 1024 * 1024 DOWNLOADED_CONTENT = 'content' diff --git a/tests/unit/s3transfer/test_crt.py b/tests/unit/s3transfer/test_crt.py index 840a39171a23..d3e5d6bf9f6f 100644 --- a/tests/unit/s3transfer/test_crt.py +++ b/tests/unit/s3transfer/test_crt.py @@ -664,6 +664,17 @@ def test_complete_is_idempotent(self): self.coordinator.complete(CustomFutureException()) assert self.coordinator.result() is None + def test_complete_releases_s3_request(self): + self.coordinator.set_s3_request(self.create_s3_request()) + self.coordinator.complete() + assert self.coordinator.s3_request is None + + def test_request_finished_before_set_is_not_kept(self): + # A request can finish before make_request() returns it. + self.coordinator.complete() + self.coordinator.set_s3_request(self.create_s3_request()) + assert self.coordinator.s3_request is None + def test_not_done_until_completed(self): s3_request = self.create_s3_request() self.coordinator.set_s3_request(s3_request)