diff --git a/README.md b/README.md index 0aa5dde0..ffd745cd 100644 --- a/README.md +++ b/README.md @@ -31,6 +31,8 @@ Profiling can be triggered in one of two ways: - `X-Profile` header format: `[=];...` +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. | | diff --git a/lib/app_profiler.rb b/lib/app_profiler.rb index 3f291b52..d3857ef1 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_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) 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..6fa6146f 100644 --- a/lib/app_profiler/middleware/upload_action.rb +++ b/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 diff --git a/lib/app_profiler/profile.rb b/lib/app_profiler/profile.rb index 376bcf45..0618e6ef 100644 --- a/lib/app_profiler/profile.rb +++ b/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 diff --git a/lib/app_profiler/railtie.rb b/lib/app_profiler/railtie.rb index 4743d9d7..89e3a96d 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.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| diff --git a/lib/app_profiler/request_parameters.rb b/lib/app_profiler/request_parameters.rb index 3c58e8f9..84f09975 100644 --- a/lib/app_profiler/request_parameters.rb +++ b/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 diff --git a/lib/app_profiler/storage/base_storage.rb b/lib/app_profiler/storage/base_storage.rb index c5309e25..23bba5f9 100644 --- a/lib/app_profiler/storage/base_storage.rb +++ b/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 diff --git a/lib/app_profiler/storage/file_storage.rb b/lib/app_profiler/storage/file_storage.rb index 129486c4..6b7334be 100644 --- a/lib/app_profiler/storage/file_storage.rb +++ b/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 diff --git a/lib/app_profiler/storage/google_cloud_storage.rb b/lib/app_profiler/storage/google_cloud_storage.rb index f4eec54b..5e184024 100644 --- a/lib/app_profiler/storage/google_cloud_storage.rb +++ b/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 diff --git a/test/app_profiler/middleware/upload_action_test.rb b/test/app_profiler/middleware/upload_action_test.rb index f6f326ca..3e188c04 100644 --- a/test/app_profiler/middleware/upload_action_test.rb +++ b/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 diff --git a/test/app_profiler/middleware_test.rb b/test/app_profiler/middleware_test.rb index 04ee0861..469fb4fd 100644 --- a/test/app_profiler/middleware_test.rb +++ b/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 diff --git a/test/app_profiler/storage/google_cloud_storage_test.rb b/test/app_profiler/storage/google_cloud_storage_test.rb index 297f400d..a0e1dc62 100644 --- a/test/app_profiler/storage/google_cloud_storage_test.rb +++ b/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