Allow out-of-band, async upload functionality

README.md

@@ -31,6 +31,8 @@ Profiling can be triggered in one of two ways:
    - `X-Profile` header format: `[<key>=<value>];...`
 
 
+If the `async` key in the query string is set to `1` or `true`, then profile will be uploaded later, in an async manner. One use case for this is when we want to profile a certain % of traffic without incurring costs of inline profile uploads. 
+
 You can configure the profile header using:
 
 ```ruby
@@ -49,6 +51,7 @@ Rails.application.config.app_profiler.profile_header = "X-Profile"
 | Key | Value | Notes |
 | --- | ----- | ----- |
 | profile/mode | Supported profiling modes: `cpu`, `wall`, `object`. | Use `profile` in (1), and `mode` in (2). |
+| async | Upload profile in a background thread. When this is set, profile redirect headers are not present in the response.
 | interval | Sampling interval in microseconds. | |
 | ignore_gc | Ignore garbage collection frames | |
 | autoredirect | Redirect request automatically to Speedscope's page after profiling. | |

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_accessor :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 :upload_queue_max_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,16 +4,19 @@ module AppProfiler
   class Middleware
     class UploadAction < BaseAction
       class << self
-        def call(profile, response: nil, autoredirect: nil)
-          profile_upload = profile.upload
-
-          return unless response
-
-          append_headers(
-            response,
-            upload: profile_upload,
-            autoredirect: autoredirect.nil? ? AppProfiler.autoredirect : autoredirect
-          )
+        def call(profile, response: nil, autoredirect: nil, async: false)
+          if async
+            profile.enqueue_upload
+            response[1][AppProfiler.profile_async_header] = true
+          else
+            profile_upload = profile.upload
+
+            append_headers(
+              response,
+              upload: profile_upload,
+              autoredirect: autoredirect.nil? ? AppProfiler.autoredirect : autoredirect
+            ) if response
+          end
         end
 
         private

lib/app_profiler/profile.rb

@@ -56,6 +56,10 @@ def upload
       nil
     end
 
+    def enqueue_upload
+      AppProfiler.storage.enqueue_upload(self)
+    end
+
     def file
       @file ||= path.tap do |p|
         p.dirname.mkpath

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.upload_queue_max_length = app.config.app_profiler.upload_queue_max_length || 10
+      AppProfiler.upload_queue_interval_secs = app.config.app_profiler.upload_queue_interval_secs || 5
     end
 
     initializer "app_profiler.add_middleware" do |app|

lib/app_profiler/request_parameters.rb

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

lib/app_profiler/storage/base_storage.rb

@@ -9,6 +9,10 @@ class BaseStorage
       def self.upload(_profile)
         raise NotImplementedError
       end
+
+      def self.enqueue_upload(_profile)
+        raise NotImplementedError
+      end
     end
   end
 end

lib/app_profiler/storage/file_storage.rb

@@ -21,6 +21,10 @@ class << self
         def upload(profile)
           Location.new(profile.file)
         end
+
+        def enqueue_upload(profile)
+          upload(profile)
+        end
       end
     end
   end

lib/app_profiler/storage/google_cloud_storage.rb

@@ -28,8 +28,59 @@ def upload(profile, _params = {})
           end
         end
 
+        def enqueue_upload(profile)
+          mutex.synchronize do
+            process_queue_thread unless @process_queue_thread&.alive?
+
+            @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
+        end
+
+        def reset_queue # for testing
+          init_queue
+          @process_queue_thread&.kill
+          @process_queue_thread = nil
+        end
+
         private
 
+        def mutex
+          @mutex ||= Mutex.new
+        end
+
+        def init_queue
+          @queue = SizedQueue.new(AppProfiler.upload_queue_max_length)
+        end
+
+        def 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
+
+        def process_queue
+          queue = nil
+          mutex.synchronize do
+            break if @queue.nil? || @queue.empty?
+
+            queue = @queue
+            init_queue
+          end
+
+          return unless queue
+
+          queue.size.times { queue.pop(false).upload }
+        end
+
         def gcs_filename(profile)
           File.join(profile.context.to_s, profile.file.basename)
         end

test/app_profiler/middleware/upload_action_test.rb

@@ -56,6 +56,8 @@ class UploadActionTest < AppProfiler::TestCase
 
         assert_predicate(@response[1][AppProfiler.profile_header], :present?)
         assert_predicate(@response[1][AppProfiler.profile_data_header], :present?)
+        assert_predicate(@response[1][AppProfiler.profile_async_header], :blank?)
+
         assert_predicate(@response[1]["Location"], :present?)
         assert_equal(@response[0], 303)
       end
@@ -81,6 +83,11 @@ class UploadActionTest < AppProfiler::TestCase
         refute_predicate(@response[1]["Location"], :present?)
       end
 
+      test ".call with async: true" do
+        UploadAction.call(@profile, response: @response, async: true)
+        assert(@response[1][AppProfiler.profile_async_header])
+      end
+
       private
 
       def with_autoredirect

test/app_profiler/middleware_test.rb

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

test/app_profiler/storage/google_cloud_storage_test.rb

@@ -9,6 +9,10 @@ class GoogleCloudStorageTest < AppProfiler::TestCase
       TEST_BUCKET_NAME = "app-profile-test"
       TEST_FILE_URL = "https://www.example.com/uploaded.json"
 
+      def teardown
+        GoogleCloudStorage.reset_queue
+      end
+
       test "upload file" do
         with_mock_gcs_bucket do
           uploaded_file = GoogleCloudStorage.upload(stub(context: "context", file: json_test_file))
@@ -46,8 +50,43 @@ class GoogleCloudStorageTest < AppProfiler::TestCase
         end
       end
 
+      test ".process_queue returns 0 when nothing to upload" do
+        Profile.any_instance.expects(:upload).never
+        GoogleCloudStorage.send(:process_queue)
+      end
+
+      test ".process_queue uploads" do
+        profile = profile_from_stackprof
+        GoogleCloudStorage.enqueue_upload(profile)
+        profile.expects(:upload).once
+        GoogleCloudStorage.send(:process_queue)
+      end
+
+      test "profile is dropped when the queue is full" do
+        AppProfiler.upload_queue_max_length.times do
+          GoogleCloudStorage.enqueue_upload(profile_from_stackprof)
+        end
+        dropped_profile = Profile.from_stackprof(profile_from_stackprof)
+        dropped_profile.expects(:upload).never
+        num_uploaded = GoogleCloudStorage.send(:process_queue)
+        assert_equal(AppProfiler.upload_queue_max_length, num_uploaded)
+      end
+
+      test "process_queue_thread is alive after first upload" do
+        th = GoogleCloudStorage.instance_variable_get(:@process_queue_thread)
+
+        refute(th&.alive?)
+        GoogleCloudStorage.enqueue_upload(profile_from_stackprof)
+        th = GoogleCloudStorage.instance_variable_get(:@process_queue_thread)
+        assert(th.alive?)
+      end
+
       private
 
+      def profile_from_stackprof
+        Profile.from_stackprof(stackprof_profile(metadata: { id: "bar" }))
+      end
+
       def json_test_file
         file_fixture("test_file.json")
       end