one enabled status

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

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

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

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