diff --git a/lib/app_profiler.rb b/lib/app_profiler.rb index 3f291b52..98898d67 100644 --- a/lib/app_profiler.rb +++ b/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) diff --git a/lib/app_profiler/middleware.rb b/lib/app_profiler/middleware.rb index 364a6b9b..cb0461f7 100644 --- a/lib/app_profiler/middleware.rb +++ b/lib/app_profiler/middleware.rb @@ -42,6 +42,7 @@ def profile(env) profile, response: response, autoredirect: params.autoredirect, + async: params.async ) response diff --git a/lib/app_profiler/middleware/upload_action.rb b/lib/app_profiler/middleware/upload_action.rb index 543dbfee..3ecde529 100644 --- a/lib/app_profiler/middleware/upload_action.rb +++ b/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 diff --git a/lib/app_profiler/railtie.rb b/lib/app_profiler/railtie.rb index 4743d9d7..d8e092e8 100644 --- a/lib/app_profiler/railtie.rb +++ b/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 diff --git a/lib/app_profiler/request_parameters.rb b/lib/app_profiler/request_parameters.rb index 3c58e8f9..bc300414 100644 --- a/lib/app_profiler/request_parameters.rb +++ b/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 diff --git a/test/app_profiler/middleware/upload_action_test.rb b/test/app_profiler/middleware/upload_action_test.rb index f6f326ca..36330d53 100644 --- a/test/app_profiler/middleware/upload_action_test.rb +++ b/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 diff --git a/test/app_profiler/middleware_test.rb b/test/app_profiler/middleware_test.rb index 04ee0861..410c1884 100644 --- a/test/app_profiler/middleware_test.rb +++ b/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