[DO NOT MERGE] TransferManager: DirectoryUploader & DirectoryDownloader

gems/aws-sdk-s3/CHANGELOG.md

@@ -1,6 +1,8 @@
 Unreleased Changes
 ------------------
 
+* Feature - TODO
+
 1.199.1 (2025-09-25)
 ------------------
 

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

@@ -6,6 +6,7 @@ module S3
     autoload :BucketRegionCache, 'aws-sdk-s3/bucket_region_cache'
     autoload :Encryption, 'aws-sdk-s3/encryption'
     autoload :EncryptionV2, 'aws-sdk-s3/encryption_v2'
+    autoload :DefaultExecutor, 'aws-sdk-s3/default_executor'
     autoload :FilePart, 'aws-sdk-s3/file_part'
     autoload :FileUploader, 'aws-sdk-s3/file_uploader'
     autoload :FileDownloader, 'aws-sdk-s3/file_downloader'
@@ -18,6 +19,9 @@ module S3
     autoload :ObjectMultipartCopier, 'aws-sdk-s3/object_multipart_copier'
     autoload :PresignedPost, 'aws-sdk-s3/presigned_post'
     autoload :Presigner, 'aws-sdk-s3/presigner'
+    autoload :DirectoryProgress, '../aws-sdk-s3/directory_progress'
+    autoload :DirectoryUploader, 'aws-sdk-s3/directory_uploader'
+    autoload :DirectoryDownloader, 'aws-sdk-s3/directory_downloader'
     autoload :TransferManager, 'aws-sdk-s3/transfer_manager'
 
     # s3 express session auth

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

@@ -459,11 +459,17 @@ def upload_stream(options = {}, &block)
       # @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)
+        executor = DefaultExecutor.new(max_threads: uploading_options.delete(:thread_count))
+        uploader = FileUploader.new(
+          client: client,
+          executor: executor,
+          multipart_threshold: uploading_options.delete(:multipart_threshold)
+        )
         response = Aws::Plugins::UserAgent.metric('RESOURCE_MODEL') do
           uploader.upload(source, uploading_options.merge(bucket: bucket_name, key: key))
         end
         yield response if block_given?
+        executor.shutdown
         true
       end
       deprecated(:upload_file, use: 'Aws::S3::TransferManager#upload_file', version: 'next major version')
@@ -539,10 +545,12 @@ def upload_file(source, options = {})
       # @see Client#get_object
       # @see Client#head_object
       def download_file(destination, options = {})
-        downloader = FileDownloader.new(client: client)
+        executor = DefaultExecutor.new(max_threads: options[: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))
         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,98 @@
+# frozen_string_literal: true
+
+module Aws
+  module S3
+    # @api private
+    class DefaultExecutor
+      RUNNING = :running
+      SHUTTING_DOWN = :shutting_down
+      SHUTDOWN = :shutdown
+
+      def initialize(options = {})
+        @max_threads = options[:max_threads] || 10
+        @state = RUNNING
+        @queue = Queue.new
+        @pool = []
+        @mutex = Mutex.new
+      end
+
+      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
+
+      def kill
+        @mutex.synchronize do
+          @state = SHUTDOWN
+          @pool.each(&:kill)
+          @pool.clear
+          @queue.clear
+        end
+        true
+      end
+
+      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
+
+        @pool.clear
+        @state = SHUTDOWN
+        true
+      end
+
+      def running?
+        @state == RUNNING
+      end
+
+      def shutting_down?
+        @state == SHUTTING_DOWN
+      end
+
+      def shutdown?
+        @state == SHUTDOWN
+      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

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

@@ -0,0 +1,175 @@
+# frozen_string_literal: true
+
+module Aws
+  module S3
+    # Raised when DirectoryDownloader fails to download objects from S3 bucket
+    class DirectoryDownloadError < StandardError
+      def initialize(message, errors = [])
+        @errors = errors
+        super(message)
+      end
+
+      # @return [Array<StandardError>] The list of errors encountered when downloading objects
+      attr_reader :errors
+    end
+
+    # @api private
+    class DirectoryDownloader
+      def initialize(options = {})
+        @client = options[:client]
+        @executor = options[:executor]
+        @abort_requested = false
+        @mutex = Mutex.new
+      end
+
+      attr_reader :abort_requested
+
+      def download(destination, bucket:, **options)
+        if File.exist?(destination)
+          raise ArgumentError, 'invalid destination, expected a directory' unless File.directory?(destination)
+        else
+          FileUtils.mkdir_p(destination)
+        end
+
+        download_opts = build_download_opts(destination, bucket, options)
+        downloader = FileDownloader.new(client: @client, executor: @executor)
+        producer = ObjectProducer.new(download_opts.merge(client: @client, directory_downloader: self))
+        downloads, errors = process_download_queue(producer, downloader, download_opts)
+        build_result(downloads, errors)
+      ensure
+        @abort_requested = false
+      end
+
+      private
+
+      def request_abort
+        @mutex.synchronize { @abort_requested = true }
+      end
+      def build_download_opts(destination, bucket, opts)
+        {
+          destination: destination,
+          bucket: bucket,
+          s3_prefix: opts.delete(:s3_prefix),
+          ignore_failure: opts.delete(:ignore_failure) || false,
+          filter_callback: opts.delete(:filter_callback),
+          progress_callback: opts.delete(:progress_callback)
+        }
+      end
+
+      def build_result(download_count, errors)
+        if @abort_requested
+          msg = "directory download failed: #{errors.map(&:message).join('; ')}"
+          raise DirectoryDownloadError.new(msg, errors)
+        else
+          {
+            completed_downloads: [download_count - errors.count, 0].max,
+            failed_downloads: errors.count,
+            errors: errors.any? ? errors : nil
+          }.compact
+        end
+      end
+
+      def handle_error(executor, opts)
+        return if opts[:ignore_failure]
+
+        request_abort
+        executor.kill
+      end
+
+      def process_download_queue(producer, downloader, opts)
+        # Separate executor for lightweight queuing tasks,
+        # avoiding interference with main @executor lifecycle
+        queue_executor = DefaultExecutor.new
+        progress = DirectoryProgress.new(opts[:progress_callback]) if opts[:progress_callback]
+        download_attempts = 0
+        errors = []
+        begin
+          producer.each do |object|
+            break if @abort_requested
+
+            download_attempts += 1
+            queue_executor.post(object) do |o|
+              dir_path = File.dirname(o[:path])
+              FileUtils.mkdir_p(dir_path) unless dir_path == opts[:destination] || Dir.exist?(dir_path)
+
+              downloader.download(o[:path], bucket: opts[:bucket], key: o[:key])
+              progress&.call(File.size(o[:path]))
+            rescue StandardError => e
+              errors << e
+              handle_error(queue_executor, opts)
+            end
+          end
+        rescue StandardError => e
+          errors << e
+          handle_error(queue_executor, opts)
+        end
+        queue_executor.shutdown
+        [download_attempts, errors]
+      end
+
+      # @api private
+      class ObjectProducer
+        include Enumerable
+
+        DEFAULT_QUEUE_SIZE = 100
+
+        def initialize(options = {})
+          @destination_dir = options[:destination]
+          @client = options[:client]
+          @bucket = options[:bucket]
+          @s3_prefix = options[:s3_prefix]
+          @filter_callback = options[:filter_callback]
+          @directory_downloader = options[:directory_downloader]
+          @object_queue = SizedQueue.new(DEFAULT_QUEUE_SIZE)
+        end
+
+        def each
+          producer_thread = Thread.new do
+            stream_objects
+          ensure
+            @object_queue << :done
+          end
+
+          # Yield objects from internal queue
+          while (object = @object_queue.shift) != :done
+            break if @directory_downloader.abort_requested
+
+            yield object
+          end
+        ensure
+          producer_thread.join
+        end
+
+        private
+
+        def build_object_entry(key)
+          { path: File.join(@destination_dir, normalize_key(key)), key: key }
+        end
+
+        # TODO: double check handling of objects that ends with /
+        def stream_objects(continuation_token: nil)
+          resp = @client.list_objects_v2(bucket: @bucket, prefix: @s3_prefix, continuation_token: continuation_token)
+          resp.contents.each do |o|
+            break if @directory_downloader.abort_requested
+            next if o.key.end_with?('/')
+            next unless include_object?(o.key)
+
+            @object_queue << build_object_entry(o.key)
+          end
+          stream_objects(continuation_token: resp.next_continuation_token) if resp.next_continuation_token
+        end
+
+        def include_object?(key)
+          return true unless @filter_callback
+
+          @filter_callback.call(key)
+        end
+
+        def normalize_key(key)
+          key = key.delete_prefix(@s3_prefix) if @s3_prefix
+          File::SEPARATOR == '/' ? key : key.tr('/', File::SEPARATOR)
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+module Aws
+  module S3
+    # @api private
+    class DirectoryProgress
+      def initialize(progress_callback)
+        @transferred_bytes = 0
+        @transferred_files = 0
+        @progress_callback = progress_callback
+        @mutex = Mutex.new
+      end
+
+      def call(bytes_transferred)
+        @mutex.synchronize do
+          @transferred_bytes += bytes_transferred
+          @transferred_files += 1
+
+          @progress_callback.call(@transferred_bytes, @transferred_files)
+        end
+      end
+    end
+  end
+end