Allow out-of-band, async upload functionality

lib/app_profiler.rb

@@ -38,6 +38,7 @@ module Viewer
   mattr_accessor :speedscope_host, default: "https://speedscope.app"
   mattr_accessor :autoredirect, default: false
   mattr_reader   :profile_header, default: "X-Profile"
+  mattr_reader   :profile_async_header, default: "X-Profile-Async"
   mattr_accessor :context, default: nil
   mattr_reader   :profile_url_formatter,
     default: DefaultProfileFormatter
@@ -46,6 +47,8 @@ module Viewer
   mattr_accessor :viewer, default: Viewer::SpeedscopeViewer
   mattr_accessor :middleware, default: Middleware
   mattr_accessor :server, default: Server
+  mattr_accessor :max_upload_queue_length, default: 10
+  mattr_accessor :upload_queue_interval_secs, default: 5
 
   class << self
     def run(*args, &block)

lib/app_profiler/middleware.rb

@@ -42,6 +42,7 @@ def profile(env)
         profile,
         response: response,
         autoredirect: params.autoredirect,
+        async: params.async
       )
 
       response

lib/app_profiler/middleware/upload_action.rb

@@ -4,9 +4,14 @@ module AppProfiler
   class Middleware
     class UploadAction < BaseAction
       class << self
-        def call(profile, response: nil, autoredirect: nil)
-          profile_upload = profile.upload
+        def call(profile, response: nil, autoredirect: nil, async: false)
+          if async
+            enqueue_upload(profile)
+            response[1][AppProfiler.profile_async_header] = true
+            return
+          end
 
+          profile_upload = profile.upload
           return unless response
 
           append_headers(
@@ -16,8 +21,43 @@ def call(profile, response: nil, autoredirect: nil)
           )
         end
 
+        def enqueue_upload(profile)
+          @queue ||= init_queue
+          begin
+            @queue.push(profile, true) # non-blocking push, raises ThreadError if queue is full
+          rescue ThreadError
+            AppProfiler.logger.info("[AppProfiler] upload queue is full, profile discarded")
+          end
+        end
+
+        def init_queue
+          @queue = SizedQueue.new(AppProfiler.max_upload_queue_length)
+        end
+
+        def start_process_queue_thread
+          @process_queue_thread ||= Thread.new do
+            loop do
+              process_queue
+              sleep(AppProfiler.upload_queue_interval_secs)
+            end
+          end
+          @process_queue_thread.priority = -1 # low priority
+        end
+
         private
 
+        def process_queue
+          return 0 if @queue.nil? || @queue.empty?
+
+          queue = @queue
+          init_queue
+
+          size = queue.length
+          size.times { queue.pop(false).upload }
+
+          size
+        end
+
         def append_headers(response, upload:, autoredirect:)
           return unless upload
 

lib/app_profiler/railtie.rb

@@ -28,11 +28,14 @@ class Railtie < Rails::Railtie
         "APP_PROFILER_SPEEDSCOPE_URL", "https://speedscope.app"
       )
       AppProfiler.profile_header = app.config.app_profiler.profile_header || "X-Profile"
+      AppProfiler.profile_async_header = app.config.app_profiler.profile_async_header || "X-Profile-Async"
       AppProfiler.profile_root = app.config.app_profiler.profile_root || Rails.root.join(
         "tmp", "app_profiler"
       )
       AppProfiler.context = app.config.app_profiler.context || Rails.env
       AppProfiler.profile_url_formatter = app.config.app_profiler.profile_url_formatter
+      AppProfiler.max_upload_queue_length = app.config.max_upload_queue_length || 10
+      AppProfiler.upload_queue_interval_secs = app.config.upload_queue_interval_secs || 5
     end
 
     initializer "app_profiler.add_middleware" do |app|
@@ -41,6 +44,9 @@ class Railtie < Rails::Railtie
           app.middleware.insert_before(0, Viewer::SpeedscopeRemoteViewer::Middleware)
         end
         app.middleware.insert_before(0, AppProfiler.middleware)
+        ActiveSupport::ForkTracker.after_fork do
+          AppProfiler::Middleware::UploadAction.start_process_queue_thread
+        end
       end
     end
 

lib/app_profiler/request_parameters.rb

@@ -16,6 +16,11 @@ def autoredirect
       query_param("autoredirect") || profile_header_param("autoredirect")
     end
 
+    def async
+      val = query_param("async")
+      val == "true" || val == "1"
+    end
+
     def valid?
       if mode.blank?
         return false

test/app_profiler/middleware/upload_action_test.rb

@@ -81,6 +81,25 @@ class UploadActionTest < AppProfiler::TestCase
         refute_predicate(@response[1]["Location"], :present?)
       end
 
+      test ".process_queue uploads" do
+        UploadAction.call(@profile, response: @response, async: true)
+        @profile.expects(:upload).once
+        assert(UploadAction.send(:process_queue) > 0)
+      end
+
+      test ".process_queue does not upload when max_upload_queue_length is exceeded" do
+        AppProfiler.max_upload_queue_length.times do
+          UploadAction.call(@profile, response: @response, async: true)
+        end
+
+        dropped_profile = Profile.new(stackprof_profile(metadata: { id: "bar" }))
+        UploadAction.call(dropped_profile, response: @response, async: true)
+
+        dropped_profile.expects(:upload).never
+        num_uploaded = UploadAction.send(:process_queue)
+        assert_equal(AppProfiler.max_upload_queue_length, num_uploaded)
+      end
+
       private
 
       def with_autoredirect

test/app_profiler/middleware_test.rb

@@ -289,6 +289,15 @@ class MiddlewareTest < TestCase
       end
     end
 
+    test "profiles are not uploaded synchronously when async is requested" do
+      assert_profiles_dumped(0) do
+        middleware = AppProfiler::Middleware.new(app_env)
+        response = middleware.call(mock_request_env(path: "/?profile=cpu&async=true"))
+        assert_equal(1, middleware.action.instance_variable_get("@queue").size)
+        assert(response[1]["X-Profile-Async"])
+      end
+    end
+
     private
 
     def app_env