lib/stackprof/middleware.rb in stackprof-0.2.6 vs lib/stackprof/middleware.rb in stackprof-0.2.7

@@ -9,28 +9,36 @@
 
       Middleware.mode     = options[:mode] || :cpu
       Middleware.interval = options[:interval] || 1000
       Middleware.enabled  = options[:enabled]
       Middleware.path     = options[:path] || 'tmp'
-      at_exit{ Middleware.save? } if options[:save_at_exit]
+      at_exit{ Middleware.save } if options[:save_at_exit]
     end
 
     def call(env)
-      StackProf.start(mode: Middleware.mode, interval: Middleware.interval) if Middleware.enabled?
+      enabled = Middleware.enabled?(env)
+      StackProf.start(mode: Middleware.mode, interval: Middleware.interval) if enabled
       @app.call(env)
     ensure
-      if Middleware.enabled?
+      if enabled
         StackProf.stop
         if @num_reqs && (@num_reqs-=1) == 0
           @num_reqs = @options[:save_every]
           Middleware.save
         end
       end
     end
 
     class << self
       attr_accessor :enabled, :mode, :interval, :path
-      alias enabled? enabled
+
+      def enabled?(env)
+        if enabled.respond_to?(:call)
+          enabled.call(env)
+        else
+          enabled
+        end
+      end
 
       def save(filename = nil)
         if results = StackProf.results
           FileUtils.mkdir_p(Middleware.path)
           filename ||= "stackprof-#{results[:mode]}-#{Process.pid}-#{Time.now.to_i}.dump"