Merge pull request #771 from flippercloud/preload-block

preload via block

lib/flipper/middleware/memoizer.rb

@@ -20,6 +20,14 @@ class Memoizer
       #   # using with preload specific features
       #   use Flipper::Middleware::Memoizer, preload: [:stats, :search, :some_feature]
       #
+      #   # using with preload block that returns true/false
+      #   use Flipper::Middleware::Memoizer, preload: ->(request) { !request.path.start_with?('/assets') }
+      #
+      #   # using with preload block that returns specific features
+      #   use Flipper::Middleware::Memoizer, preload: ->(request) {
+      #     request.path.starts_with?('/admin') ? [:stats, :search] : false
+      #   }
+      #
       def initialize(app, opts = {})
         if opts.is_a?(Flipper::DSL) || opts.is_a?(Proc)
           raise 'Flipper::Middleware::Memoizer no longer initializes with a flipper instance or block. Read more at: https://git.io/vSo31.'
@@ -34,7 +42,7 @@ def call(env)
         request = Rack::Request.new(env)
 
         if memoize?(request)
-          memoized_call(env)
+          memoized_call(request)
         else
           @app.call(env)
         end
@@ -52,26 +60,34 @@ def memoize?(request)
         end
       end
 
-      def memoized_call(env)
-        reset_on_body_close = false
-        flipper = env.fetch(@env_key) { Flipper }
+      def memoized_call(request)
+        flipper = request.env.fetch(@env_key) { Flipper }
 
         # Already memoizing. This instance does not need to do anything.
         if flipper.memoizing?
           warn "Flipper::Middleware::Memoizer appears to be running twice. Read how to resolve this at https://github.com/flippercloud/flipper/pull/523"
-          return @app.call(env)
+          return @app.call(request.env)
         end
 
-        flipper.memoize = true
+        begin
+          flipper.memoize = true
 
-        case @opts[:preload]
-        when true then flipper.preload_all
-        when Array then flipper.preload(@opts[:preload])
-        end
+          # Preloading is pointless without memoizing.
+          preload = if @opts[:preload].respond_to?(:call)
+            @opts[:preload].call(request)
+          else
+            @opts[:preload]
+          end
 
-        @app.call(env)
-      ensure
-        flipper.memoize = false if flipper
+          case preload
+          when true then flipper.preload_all
+          when Array then flipper.preload(preload)
+          end
+
+          @app.call(request.env)
+        ensure
+          flipper.memoize = false
+        end
       end
     end
   end

spec/flipper/middleware/memoizer_spec.rb

@@ -196,6 +196,73 @@
     end
   end
 
+  context 'with preload block' do
+    let(:app) do
+      app = lambda do |_env|
+        flipper[:stats].enabled?
+        flipper[:stats].enabled?
+        flipper[:shiny].enabled?
+        flipper[:shiny].enabled?
+        [200, {}, []]
+      end
+
+      described_class.new(app, preload: ->(request) {
+        case request.path
+        when "/true"
+          true
+        when "/specific"
+          [:stats]
+        else
+          false
+        end
+      })
+    end
+
+    include_examples 'flipper middleware'
+
+    it 'eagerly caches known features for duration of request if block returns true' do
+      flipper[:stats].enable
+      flipper[:shiny].enable
+
+      # clear the log of operations
+      adapter.reset
+
+      get '/true', {}, 'flipper' => flipper
+
+      expect(adapter.operations.size).to be(1)
+      expect(adapter.count(:get_all)).to be(1)
+      expect(adapter.count(:get)).to be(0)
+    end
+
+    it 'does not eagerly cache known features if block returns false' do
+      flipper[:stats].enable
+      flipper[:shiny].enable
+
+      # clear the log of operations
+      adapter.reset
+
+      get '/false', {}, 'flipper' => flipper
+
+      expect(adapter.operations.size).to be(2)
+      expect(adapter.count(:get_all)).to be(0)
+      expect(adapter.count(:get)).to be(2)
+    end
+
+    it 'eagerly caches specified features for duration of request if block returns array of specified features' do
+      flipper[:stats].enable
+      flipper[:shiny].enable
+
+      # clear the log of operations
+      adapter.reset
+
+      get '/specific', {}, 'flipper' => flipper
+
+      expect(adapter.operations.size).to be(2)
+      expect(adapter.count(:get_multi)).to be(1)
+      expect(adapter.count(:get)).to be(1)
+    end
+  end
+
   context 'with multiple instances' do
     let(:app) do
       # ensure scoped for builder block, annoying...