S3 Executor support (#3302)

Author: jterapin

gems/aws-sdk-s3/CHANGELOG.md

@@ -1,6 +1,10 @@
 Unreleased Changes
 ------------------
 
+* Feature - Add lightweight thread pool executor for multipart `download_file`, `upload_file` and `upload_stream`.
+
+* Feature - Add custom executor support for `Aws::S3::TransferManager`.
+
 1.199.1 (2025-09-25)
 ------------------
 

gems/aws-sdk-s3/lib/aws-sdk-s3/customizations.rb

@@ -7,6 +7,7 @@ module S3
     autoload :Encryption, 'aws-sdk-s3/encryption'
     autoload :EncryptionV2, 'aws-sdk-s3/encryption_v2'
     autoload :FilePart, 'aws-sdk-s3/file_part'
+    autoload :DefaultExecutor, 'aws-sdk-s3/default_executor'
     autoload :FileUploader, 'aws-sdk-s3/file_uploader'
     autoload :FileDownloader, 'aws-sdk-s3/file_downloader'
     autoload :LegacySigner, 'aws-sdk-s3/legacy_signer'

gems/aws-sdk-s3/lib/aws-sdk-s3/customizations/object.rb

@@ -358,8 +358,8 @@ def public_url(options = {})
       #   {Client#complete_multipart_upload},
       #   and {Client#upload_part} can be provided.
       #
-      # @option options [Integer] :thread_count (10) The number of parallel
-      #   multipart uploads
+      # @option options [Integer] :thread_count (10) The number of parallel multipart uploads.
+      #   An additional thread is used internally for task coordination.
       #
       # @option options [Boolean] :tempfile (false) Normally read data is stored
       #   in memory when building the parts in order to complete the underlying
@@ -383,19 +383,18 @@ def public_url(options = {})
       # @see Client#complete_multipart_upload
       # @see Client#upload_part
       def upload_stream(options = {}, &block)
-        uploading_options = options.dup
+        upload_opts = options.merge(bucket: bucket_name, key: key)
+        executor = DefaultExecutor.new(max_threads: upload_opts.delete(:thread_count))
         uploader = MultipartStreamUploader.new(
           client: client,
-          thread_count: uploading_options.delete(:thread_count),
-          tempfile: uploading_options.delete(:tempfile),
-          part_size: uploading_options.delete(:part_size)
+          executor: executor,
+          tempfile: upload_opts.delete(:tempfile),
+          part_size: upload_opts.delete(:part_size)
         )
         Aws::Plugins::UserAgent.metric('RESOURCE_MODEL') do
-          uploader.upload(
-            uploading_options.merge(bucket: bucket_name, key: key),
-            &block
-          )
+          uploader.upload(upload_opts, &block)
         end
+        executor.shutdown
         true
       end
       deprecated(:upload_stream, use: 'Aws::S3::TransferManager#upload_stream', version: 'next major version')
@@ -458,12 +457,18 @@ def upload_stream(options = {}, &block)
       # @see Client#complete_multipart_upload
       # @see Client#upload_part
       def upload_file(source, options = {})
-        uploading_options = options.dup
-        uploader = FileUploader.new(multipart_threshold: uploading_options.delete(:multipart_threshold), client: client)
+        upload_opts = options.merge(bucket: bucket_name, key: key)
+        executor = DefaultExecutor.new(max_threads: upload_opts.delete(:thread_count))
+        uploader = FileUploader.new(
+          client: client,
+          executor: executor,
+          multipart_threshold: upload_opts.delete(:multipart_threshold)
+        )
         response = Aws::Plugins::UserAgent.metric('RESOURCE_MODEL') do
-          uploader.upload(source, uploading_options.merge(bucket: bucket_name, key: key))
+          uploader.upload(source, upload_opts)
         end
         yield response if block_given?
+        executor.shutdown
         true
       end
       deprecated(:upload_file, use: 'Aws::S3::TransferManager#upload_file', version: 'next major version')
@@ -512,10 +517,6 @@ def upload_file(source, options = {})
       #
       # @option options [Integer] :thread_count (10) Customize threads used in the multipart download.
       #
-      # @option options [String] :version_id The object version id used to retrieve the object.
-      #
-      #     @see https://docs.aws.amazon.com/AmazonS3/latest/dev/ObjectVersioning.html ObjectVersioning
-      #
       # @option options [String] :checksum_mode ("ENABLED")
       #   When `"ENABLED"` and the object has a stored checksum, it will be used to validate the download and will
       #   raise an `Aws::Errors::ChecksumError` if checksum validation fails. You may provide a `on_checksum_validated`
@@ -539,10 +540,13 @@ def upload_file(source, options = {})
       # @see Client#get_object
       # @see Client#head_object
       def download_file(destination, options = {})
-        downloader = FileDownloader.new(client: client)
+        download_opts = options.merge(bucket: bucket_name, key: key)
+        executor = DefaultExecutor.new(max_threads: download_opts.delete([:thread_count]))
+        downloader = FileDownloader.new(client: client, executor: executor)
         Aws::Plugins::UserAgent.metric('RESOURCE_MODEL') do
-          downloader.download(destination, options.merge(bucket: bucket_name, key: key))
+          downloader.download(destination, download_opts)
         end
+        executor.shutdown
         true
       end
       deprecated(:download_file, use: 'Aws::S3::TransferManager#download_file', version: 'next major version')

gems/aws-sdk-s3/lib/aws-sdk-s3/default_executor.rb

@@ -0,0 +1,103 @@
+# frozen_string_literal: true
+
+module Aws
+  module S3
+    # @api private
+    class DefaultExecutor
+      DEFAULT_MAX_THREADS = 10
+      RUNNING = :running
+      SHUTTING_DOWN = :shutting_down
+      SHUTDOWN = :shutdown
+
+      def initialize(options = {})
+        @max_threads = options[:max_threads] || DEFAULT_MAX_THREADS
+        @state = RUNNING
+        @queue = Queue.new
+        @pool = []
+        @mutex = Mutex.new
+      end
+
+      # Submits a task for execution.
+      # @param [Object] args Variable number of arguments to pass to the block
+      # @param [Proc] block The block to be executed
+      # @return [Boolean] Returns true if the task was submitted successfully
+      def post(*args, &block)
+        @mutex.synchronize do
+          raise 'Executor has been shutdown and is no longer accepting tasks' unless @state == RUNNING
+
+          @queue << [args, block]
+          ensure_worker_available
+        end
+        true
+      end
+
+      # Immediately terminates all worker threads and clears pending tasks.
+      # This is a forceful shutdown that doesn't wait for running tasks to complete.
+      #
+      # @return [Boolean] true when termination is complete
+      def kill
+        @mutex.synchronize do
+          @state = SHUTDOWN
+          @pool.each(&:kill)
+          @pool.clear
+          @queue.clear
+        end
+        true
+      end
+
+      # Gracefully shuts down the executor, optionally with a timeout.
+      # Stops accepting new tasks and waits for running tasks to complete.
+      #
+      # @param timeout [Numeric, nil] Maximum time in seconds to wait for shutdown.
+      #   If nil, waits indefinitely. If timeout expires, remaining threads are killed.
+      # @return [Boolean] true when shutdown is complete
+      def shutdown(timeout = nil)
+        @mutex.synchronize do
+          return true if @state == SHUTDOWN
+
+          @state = SHUTTING_DOWN
+          @pool.size.times { @queue << :shutdown }
+        end
+
+        if timeout
+          deadline = Time.now + timeout
+          @pool.each do |thread|
+            remaining = deadline - Time.now
+            break if remaining <= 0
+
+            thread.join([remaining, 0].max)
+          end
+          @pool.select(&:alive?).each(&:kill)
+        else
+          @pool.each(&:join)
+        end
+
+        @mutex.synchronize do
+          @pool.clear
+          @state = SHUTDOWN
+        end
+        true
+      end
+
+      private
+
+      def ensure_worker_available
+        return unless @state == RUNNING
+
+        @pool.select!(&:alive?)
+        @pool << spawn_worker if @pool.size < @max_threads
+      end
+
+      def spawn_worker
+        Thread.new do
+          while (job = @queue.shift)
+            break if job == :shutdown
+
+            args, block = job
+            block.call(*args)
+          end
+        end
+      end
+    end
+  end
+end