diff --git a/lib/app_profiler.rb b/lib/app_profiler.rb index 3006efe..00f5ca2 100644 --- a/lib/app_profiler.rb +++ b/lib/app_profiler.rb @@ -63,8 +63,7 @@ module Viewer mattr_reader :profile_enqueue_failure, default: nil mattr_reader :after_process_queue, default: nil mattr_accessor :forward_metadata_on_upload, default: false - mattr_accessor :profile_sampler_enabled, default: false - mattr_accessor :profile_sampler_killswitch_on, default: nil + mattr_accessor :profile_sampler_config class << self @@ -117,6 +116,25 @@ def backend=(new_backend) @profiler_backend = new_profiler_backend end + def profile_sampler_enabled=(value) + if value.is_a?(Proc) + raise ArgumentError, + "profile_sampler_enabled must be a proc or a lambda that accepts no argument" if value.arity != 0 + else + raise ArgumentError, "Must be TrueClass or FalseClass" unless [TrueClass, FalseClass].include?(value.class) + end + + @profile_sampler_enabled = value + end + + def profile_sampler_enabled + return false unless defined?(@profile_sampler_enabled) + + @profile_sampler_enabled.is_a?(Proc) ? @profile_sampler_enabled.call : @profile_sampler_enabled + rescue + false + end + def backend_for(backend_name) if vernier_supported? && backend_name == AppProfiler::Backend::VernierBackend.name diff --git a/lib/app_profiler/middleware.rb b/lib/app_profiler/middleware.rb index 053e927..bd5a63e 100644 --- a/lib/app_profiler/middleware.rb +++ b/lib/app_profiler/middleware.rb @@ -50,9 +50,7 @@ def profile(env, params) def profile_params(params) return params if params.valid? - return unless AppProfiler.profile_sampler_enabled - return if AppProfiler.profile_sampler_killswitch_on&.call AppProfiler::Sampler.profile_params(params, AppProfiler.profile_sampler_config) end diff --git a/lib/app_profiler/railtie.rb b/lib/app_profiler/railtie.rb index 5c9f0bc..7e616b9 100644 --- a/lib/app_profiler/railtie.rb +++ b/lib/app_profiler/railtie.rb @@ -43,7 +43,6 @@ class Railtie < Rails::Railtie AppProfiler.backend = app.config.app_profiler.profiler_backend || :stackprof AppProfiler.forward_metadata_on_upload = app.config.app_profiler.forward_metadata_on_upload || false AppProfiler.profile_sampler_enabled = app.config.app_profiler.profile_sampler_enabled || false - AppProfiler.profile_sampler_killswitch_on = app.config.app_profiler.profile_sampler_killswitch_on || nil AppProfiler.profile_sampler_config = app.config.app_profiler.profile_sampler_config || AppProfiler::Sampler::Config.new end diff --git a/test/app_profiler/middleware_test.rb b/test/app_profiler/middleware_test.rb index 1159be3..05703d4 100644 --- a/test/app_profiler/middleware_test.rb +++ b/test/app_profiler/middleware_test.rb @@ -399,19 +399,16 @@ def call(env) end end - test "request is not sampled when killswitch is on" do - old_killswitch = AppProfiler.profile_sampler_killswitch_on + test "request is not sampled when sampler is not enabled via Proc" do with_google_cloud_storage do AppProfiler.profile_sampler_config = AppProfiler::Sampler::Config.new(sample_rate: 1.0) - AppProfiler.profile_sampler_killswitch_on = -> { true } + AppProfiler.profile_sampler_enabled = -> { false } assert_profiles_dumped(0) do middleware = AppProfiler::Middleware.new(app_env) response = middleware.call(mock_request_env(path: "/")) assert_nil(response[1]["X-Profile-Async"]) end end - ensure - AppProfiler.profile_sampler_killswitch_on = old_killswitch end private @@ -421,6 +418,7 @@ def with_profile_sampler_enabled AppProfiler.profile_sampler_enabled = true yield ensure + puts "status #{old_status}" AppProfiler.profile_sampler_enabled = old_status end