From d5a2f7ec148726d7547e367d7a968e3b4be9b509 Mon Sep 17 00:00:00 2001
From: George Claghorn <george@basecamp.com>
Date: Wed, 22 May 2019 15:07:35 -0400
Subject: Mirror direct uploads

---
 activestorage/lib/active_storage/engine.rb         |  4 +--
 activestorage/lib/active_storage/log_subscriber.rb |  6 +++++
 .../lib/active_storage/service/mirror_service.rb   | 29 +++++++++++++++++++---
 3 files changed, 33 insertions(+), 6 deletions(-)

(limited to 'activestorage/lib')

diff --git a/activestorage/lib/active_storage/engine.rb b/activestorage/lib/active_storage/engine.rb
index cbb205627e..f70d0a512a 100644
--- a/activestorage/lib/active_storage/engine.rb
+++ b/activestorage/lib/active_storage/engine.rb
@@ -24,7 +24,7 @@ module ActiveStorage
     config.active_storage.previewers = [ ActiveStorage::Previewer::PopplerPDFPreviewer, ActiveStorage::Previewer::MuPDFPreviewer, ActiveStorage::Previewer::VideoPreviewer ]
     config.active_storage.analyzers = [ ActiveStorage::Analyzer::ImageAnalyzer, ActiveStorage::Analyzer::VideoAnalyzer ]
     config.active_storage.paths = ActiveSupport::OrderedOptions.new
-    config.active_storage.queues = ActiveSupport::OrderedOptions.new
+    config.active_storage.queues = ActiveSupport::InheritableOptions.new(mirror: :active_storage_mirror)
 
     config.active_storage.variable_content_types = %w(
       image/png
@@ -130,7 +130,7 @@ module ActiveStorage
             "config.active_storage.queue is deprecated and will be removed in Rails 6.1. " \
             "Set config.active_storage.queues.purge and config.active_storage.queues.analysis instead."
 
-          ActiveStorage.queues = { purge: queue, analysis: queue }
+          ActiveStorage.queues = { purge: queue, analysis: queue, mirror: queue }
         else
           ActiveStorage.queues = app.config.active_storage.queues || {}
         end
diff --git a/activestorage/lib/active_storage/log_subscriber.rb b/activestorage/lib/active_storage/log_subscriber.rb
index 6c0b4c30e7..f823f31093 100644
--- a/activestorage/lib/active_storage/log_subscriber.rb
+++ b/activestorage/lib/active_storage/log_subscriber.rb
@@ -32,6 +32,12 @@ module ActiveStorage
       debug event, color("Generated URL for file at key: #{key_in(event)} (#{event.payload[:url]})", BLUE)
     end
 
+    def service_mirror(event)
+      message = "Mirrored file at key: #{key_in(event)}"
+      message += " (checksum: #{event.payload[:checksum]})" if event.payload[:checksum]
+      debug event, color(message, GREEN)
+    end
+
     def logger
       ActiveStorage.logger
     end
diff --git a/activestorage/lib/active_storage/service/mirror_service.rb b/activestorage/lib/active_storage/service/mirror_service.rb
index aa41df304e..c44bd1f360 100644
--- a/activestorage/lib/active_storage/service/mirror_service.rb
+++ b/activestorage/lib/active_storage/service/mirror_service.rb
@@ -4,12 +4,17 @@ require "active_support/core_ext/module/delegation"
 
 module ActiveStorage
   # Wraps a set of mirror services and provides a single ActiveStorage::Service object that will all
-  # have the files uploaded to them. A +primary+ service is designated to answer calls to +download+, +exists?+,
-  # and +url+.
+  # have the files uploaded to them. A +primary+ service is designated to answer calls to:
+  # * +download+
+  # * +exists?+
+  # * +url+
+  # * +url_for_direct_upload+
+  # * +headers_for_direct_upload+
   class Service::MirrorService < Service
     attr_reader :primary, :mirrors
 
-    delegate :download, :download_chunk, :exist?, :url, :path_for, to: :primary
+    delegate :download, :download_chunk, :exist?, :url,
+      :url_for_direct_upload, :headers_for_direct_upload, :path_for, to: :primary
 
     # Stitch together from named services.
     def self.build(primary:, mirrors:, configurator:, **options) #:nodoc:
@@ -26,7 +31,8 @@ module ActiveStorage
     # ensure a match when the upload has completed or raise an ActiveStorage::IntegrityError.
     def upload(key, io, checksum: nil, **options)
       each_service.collect do |service|
-        service.upload key, io.tap(&:rewind), checksum: checksum, **options
+        io.rewind
+        service.upload key, io, checksum: checksum, **options
       end
     end
 
@@ -40,6 +46,21 @@ module ActiveStorage
       perform_across_services :delete_prefixed, prefix
     end
 
+
+    # Copy the file at the +key+ from the primary service to each of the mirrors where it doesn't already exist.
+    def mirror(key, checksum:)
+      instrument :mirror, key: key, checksum: checksum do
+        if (mirrors_in_need_of_mirroring = mirrors.select { |service| !service.exist?(key) }).any?
+          primary.open(key, checksum: checksum) do |io|
+            mirrors_in_need_of_mirroring.each do |service|
+              io.rewind
+              service.upload key, io, checksum: checksum
+            end
+          end
+        end
+      end
+    end
+
     private
       def each_service(&block)
         [ primary, *mirrors ].each(&block)
-- 
cgit v1.2.3