lib/flipper/middleware/memoizer.rb in flipper-0.11.0.beta7 vs lib/flipper/middleware/memoizer.rb in flipper-0.11.0.beta8

@@ -1,13 +1,14 @@
 require 'rack/body_proxy'
 
 module Flipper
   module Middleware
     class Memoizer
-      # Public: Initializes an instance of the Memoizer middleware. The flipper
-      # instance must be setup in the env of the request. You can do this by
-      # using the Flipper::Middleware::SetupEnv middleware.
+      # Public: Initializes an instance of the Memoizer middleware. Flipper must
+      # be configured with a default instance or the flipper instance must be
+      # setup in the env of the request. You can do this by using the
+      # Flipper::Middleware::SetupEnv middleware.
       #
       # app - The app this middleware is included in.
       # opts - The Hash of options.
       #        :preload_all - Boolean of whether or not to preload all features.
       #        :preload - Array of Symbol feature names to preload.
@@ -47,11 +48,12 @@
       def skip_memoize?(request)
         @opts[:unless] && @opts[:unless].call(request)
       end
 
       def memoized_call(env)
-        flipper = env.fetch(@env_key)
+        reset_on_body_close = false
+        flipper = env.fetch(@env_key) { Flipper }
         original = flipper.adapter.memoizing?
         flipper.adapter.memoize = true
 
         flipper.preload_all if @opts[:preload_all]
 
@@ -61,13 +63,13 @@
 
         response = @app.call(env)
         response[2] = Rack::BodyProxy.new(response[2]) do
           flipper.adapter.memoize = original
         end
+        reset_on_body_close = true
         response
-      rescue
-        flipper.adapter.memoize = original
-        raise
+      ensure
+        flipper.adapter.memoize = original if flipper && !reset_on_body_close
       end
     end
   end
 end