[Sampler] Allow passing in a Proc to enable/disable the sampler.

README.md

@@ -341,22 +341,29 @@ By default, the stackprof backend will be used.
 
 The `AppProfiler` middleware can be configured to sample a percentage of requests for profiling. This can be useful for profiling a subset of requests in production without profiling every request.
 
-To enable profile sampling:
+To enable profile sampling, you can either set `profile_sampler_enabled` to `true` or provide your own custom `Proc`:
 
 ```ruby
 AppProfiler.profile_sampler_enabled = true
 # OR
 Rails.application.config.app_profiler.profile_sampler_enabled = true
 ```
 
+```ruby
+AppProfiler.profile_sampler_enabled = -> { sampler_enabled? }
+# OR
+Rails.application.config.app_profiler.profile_sampler_enabled = -> { sampler_enabled? }
+```
+
+
 Both backends, StackProf and Vernier, can be configured separately.
 
 These can be overridden like:
 
 ```ruby
 AppProfiler.profile_sampler_config = AppProfiler::Sampler::Config.new(
   sample_rate: 0.5,
-  paths: ["/foo"],
+  targets: ["/foo"],
   backends_probability: { stackprof: 0.5, vernier: 0.5 },
   backend_configs: {
     stackprof: AppProfiler::Sampler::StackprofConfig.new,
@@ -367,7 +374,7 @@ AppProfiler.profile_sampler_config = AppProfiler::Sampler::Config.new(
 
 Rails.application.config.app_profiler = AppProfiler::Sampler::Config.new(
   sample_rate: 0.5,
-  paths: ["/foo"],
+  targets: ["/foo"],
   backends_probability: { stackprof: 0.5, vernier: 0.5 },
   backend_configs: {
     stackprof: AppProfiler::Sampler::StackprofConfig.new,
@@ -380,12 +387,12 @@ All the configuration parameters are optional and have default values. The defau
 
 ```ruby
 
-| Sampler Config                      | Default         |
-| --------                            | -------         |
-| sample_rate (0.0 - 1.0)             | 0.001 (0.1 %)   |
-| paths                               | ['/']           |
-| stackprof_probability               | 1.0             |
-| vernier_probability                 | 0.0             |
+| Sampler Config                                         | Default         |
+| --------                                               | -------         |
+| sample_rate (0.0 - 1.0)                                | 0.001 (0.1 %)   |
+| targets (request paths, job_names etc )                | ['/']           |
+| stackprof_probability                                  | 1.0             |
+| vernier_probability                                    | 0.0             |
 
 | StackprofConfig                     | Default |
 | --------                            | ------- |

lib/app_profiler.rb

@@ -63,7 +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_config
 
   class << self
@@ -116,6 +116,28 @@ 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 => e
+      logger.error(
+        "[AppProfiler.profile_sampler_enabled] exception: #{e}, message: #{e.message}",
+      )
+      false
+    end
+
     def backend_for(backend_name)
       if vernier_supported? &&
           backend_name == AppProfiler::Backend::VernierBackend.name

lib/app_profiler/middleware.rb

@@ -50,7 +50,6 @@ def profile(env, params)
 
     def profile_params(params)
       return params if params.valid?
-
       return unless AppProfiler.profile_sampler_enabled
 
       AppProfiler::Sampler.profile_params(params, AppProfiler.profile_sampler_config)

lib/app_profiler/sampler.rb

@@ -5,18 +5,21 @@ module AppProfiler
   module Sampler
     class << self
       def profile_params(request, config)
-        return unless sample?(config, request)
+        profile_params_for(request.path, config)
+      end
+
+      def profile_params_for(target, config)
+        return unless sample?(config, target)
 
         get_profile_params(config)
       end
 
       private
 
-      def sample?(config, request)
+      def sample?(config, target)
         return false if Kernel.rand > config.sample_rate
 
-        path = request.path
-        return false unless config.paths.any? { |p| path.match?(p) }
+        return false unless config.targets.any? { |t| target.match?(t) }
 
         true
       end

lib/app_profiler/sampler/config.rb

@@ -5,28 +5,31 @@
 module AppProfiler
   module Sampler
     class Config
-      attr_reader :sample_rate, :paths, :cpu_interval, :backends_probability
+      attr_reader :sample_rate, :targets, :cpu_interval, :backends_probability
 
       SAMPLE_RATE = 0.001 # 0.1%
-      PATHS = ["/"]
+      TARGETS = ["/"]
       BACKEND_PROBABILITES = { stackprof: 1.0, vernier: 0.0 }
       @backends = {}
 
       def initialize(sample_rate: SAMPLE_RATE,
-        paths: PATHS,
+        targets: TARGETS,
         backends_probability: BACKEND_PROBABILITES,
         backends_config: {
           stackprof: StackprofConfig.new,
-        })
+        },
+        paths: nil)
 
         if sample_rate < 0.0 || sample_rate > 1.0
           raise ArgumentError, "sample_rate must be between 0 and 1"
         end
 
         raise ArgumentError, "mode probabilities must sum to 1" unless backends_probability.values.sum == 1.0
 
+        ActiveSupport::Deprecation.new.warn("passing paths is deprecated, use targets instead") if paths
+
         @sample_rate = sample_rate
-        @paths = paths
+        @targets = paths || targets
         @backends_config = backends_config
         @backends_probability = backends_probability
       end

test/app_profiler/app_profiler_configuration_test.rb

@@ -20,5 +20,34 @@ class ConfigurationTest < TestCase
         AppProfiler.after_process_queue = lambda { nil }
       end
     end
+
+    test "unexpected handler profile_sampler_enabled raises ArgumentError " do
+      assert_raises(ArgumentError) do
+        AppProfiler.profile_sampler_enabled = ->(arg) { arg }
+      end
+    end
+
+    test "proc with args for profile_sampler_enabled raises ArgumentError " do
+      assert_raises(ArgumentError) do
+        AppProfiler.profile_sampler_enabled = ->(arg) { arg }
+      end
+    end
+
+    test "profile_sampler_enabled allows TrueClass, FalseClass, Proc" do
+      old_status = AppProfiler.profile_sampler_enabled
+      AppProfiler.profile_sampler_enabled = false
+      AppProfiler.profile_sampler_enabled = true
+      AppProfiler.profile_sampler_enabled = -> { false }
+    ensure
+      AppProfiler.profile_sampler_enabled = old_status
+    end
+
+    test "profile_sampler_enabled is false if Proc raises" do
+      old_status = AppProfiler.profile_sampler_enabled
+      AppProfiler.profile_sampler_enabled = -> { raise "error" }
+      assert_equal(false, AppProfiler.profile_sampler_enabled)
+    ensure
+      AppProfiler.profile_sampler_enabled = old_status
+    end
   end
 end

test/app_profiler/middleware_test.rb

@@ -376,7 +376,7 @@ def call(env)
         with_google_cloud_storage do
           AppProfiler.profile_sampler_config = AppProfiler::Sampler::Config.new(
             sample_rate: 1.0,
-            paths: ["/"],
+            targets: ["/"],
           )
 
           assert_profiles_dumped(0) do
@@ -399,6 +399,18 @@ def call(env)
       end
     end
 
+    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_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
+    end
+
     private
 
     def with_profile_sampler_enabled

test/app_profiler/sampler/config_test.rb

@@ -21,11 +21,16 @@ class ConfigTest < TestCase
       test "default config" do
         config = Config.new
         assert_equal(Config::SAMPLE_RATE, config.sample_rate)
-        assert_equal(Config::PATHS, config.paths)
+        assert_equal(Config::TARGETS, config.targets)
         assert_equal(Config::BACKEND_PROBABILITES, config.backends_probability)
         assert_not_nil(config.get_backend_config(:stackprof))
         assert_nil(config.get_backend_config(:vernier))
       end
+
+      test "allows passing paths" do
+        config = Config.new(paths: ["/"])
+        assert_equal(["/"], config.targets)
+      end
     end
   end
 end

test/app_profiler/sampler/sampler_test.rb

@@ -48,7 +48,7 @@ class SamplerTest < TestCase
         Kernel.stubs(:rand).returns(0.1)
         config = Config.new(
           sample_rate: 1.0,
-          paths: ["/foo"],
+          targets: ["/foo"],
         )
 
         request = RequestParameters.new(Rack::Request.new({ "PATH_INFO" => "/foo" }))