Use fetch_multi for multiple cache blocks

lib/jbuilder.rb

@@ -1,4 +1,5 @@
 require 'jbuilder/jbuilder'
+require 'jbuilder/cache'
 require 'jbuilder/blank'
 require 'jbuilder/key_formatter'
 require 'jbuilder/errors'

lib/jbuilder/cache.rb

@@ -0,0 +1,37 @@
+class Jbuilder
+  class Cache
+    def initialize
+      @cache = Hash.new { |h, k| h[k] = {} }
+    end
+
+    def add(key, options, &block)
+      @cache[options][key] = block
+    end
+
+    def resolve
+      resolved = []
+
+      # Fail-fast if there is no items to be computed.
+      return resolved if @cache.empty?
+
+      # We can't add new items during interation, so iterate through a
+      # clone that will allow us to add new items.
+      cache = @cache.clone
+      @cache.clear
+
+      # Keys are grouped by options and because of that, fetch_multi
+      # will use the same options for the same group of keys.
+      cache.each do |options, group|
+        result = Rails.cache.fetch_multi(*group.keys, options) do |group_key|
+          [group[group_key].call, *resolve]
+        end
+
+        # Since the fetch_multi returns { cache_key => value }, we need
+        # to discard the cache key and merge only the values.
+        resolved.concat result.values.flatten(1)
+      end
+
+      resolved
+    end
+  end
+end

lib/jbuilder/jbuilder_template.rb

@@ -12,6 +12,7 @@ class << self
   def initialize(context, *args)
     @context = context
     @cached_root = nil
+    @cache = Cache.new
     super(*args)
   end
 
@@ -33,11 +34,9 @@ def partial!(*args)
   #   end
   def cache!(key=nil, options={})
     if @context.controller.perform_caching
-      value = _cache_fragment_for(key, options) do
+      _cache_fragment_for(key, options) do
         _scope { yield self }
       end
-
-      merge! value
     else
       yield
     end
@@ -58,7 +57,8 @@ def cache_root!(key=nil, options={})
     if @context.controller.perform_caching
       raise "cache_root! can't be used after JSON structures have been defined" if @attributes.present?
 
-      @cached_root = _cache_fragment_for([ :root, key ], options) { yield; target! }
+      cache_key = _cache_key([ :root, key ], options)
+      @cached_root = ::Rails.cache.fetch(cache_key, options) { yield; target! }
     else
       yield
     end
@@ -78,7 +78,13 @@ def cache_if!(condition, *args)
   end
 
   def target!
-    @cached_root || super
+    return @cached_root if @cached_root
+
+    @cache.resolve.each do |value|
+      @attributes = _merge_values(@attributes, value)
+    end
+
+    super
   end
 
   def array!(collection = [], *args)
@@ -130,21 +136,7 @@ def _render_partial(options)
 
   def _cache_fragment_for(key, options, &block)
     key = _cache_key(key, options)
-    _read_fragment_cache(key, options) || _write_fragment_cache(key, options, &block)
-  end
-
-  def _read_fragment_cache(key, options = nil)
-    @context.controller.instrument_fragment_cache :read_fragment, key do
-      ::Rails.cache.read(key, options)
-    end
-  end
-
-  def _write_fragment_cache(key, options = nil)
-    @context.controller.instrument_fragment_cache :write_fragment, key do
-      yield.tap do |value|
-        ::Rails.cache.write(key, value, options)
-      end
-    end
+    @cache.add(key, options, &block)
   end
 
   def _cache_key(key, options)

test/jbuilder_cache_test.rb

@@ -0,0 +1,132 @@
+require "test_helper"
+require "jbuilder"
+require "jbuilder/cache"
+
+class JbuilderCacheTest < ActiveSupport::TestCase
+  setup do
+    Rails.cache.clear
+  end
+
+  test "resolve flat" do
+    cache = Jbuilder::Cache.new
+    cache.add("x", {}) { "x" }
+    cache.add("y", {}) { [ "y" ] }
+    cache.add("z", {}) { { key: "value" } }
+
+    assert_equal(["x", ["y"], { key: "value" }], cache.resolve)
+  end
+
+  test "resolve nested" do
+    cache = Jbuilder::Cache.new
+    cache.add("x", {}) do
+      cache.add("y", {}) do
+        cache.add("z", {}) do
+          { key: "value" }
+        end
+
+        ["y"]
+      end
+
+      "x"
+    end
+
+    assert_equal(["x", ["y"], { key: "value" }], cache.resolve)
+  end
+
+  test "cache calls" do
+    3.times do
+      cache = Jbuilder::Cache.new
+      cache.add("x", {}) { "x" }
+      cache.add("y", {}) { [ "y" ] }
+      cache.add("z", { expires_in: 10.minutes }) { { key: "value" } }
+      cache.resolve
+    end
+
+    # cache miss:
+    #
+    #   1. x + y
+    #   2. z
+    #
+    # cache hit:
+    #
+    #   3. x + y
+    #   4. z
+    #   5. x + y
+    #   6. z
+    assert_equal 6, Rails.cache.fetch_multi_calls.length
+
+    # cache miss:
+    #
+    #   1. x
+    #   2. y
+    #   2. z
+    assert_equal 3, Rails.cache.write_calls.length
+  end
+
+  test "nested cache calls" do
+    3.times do
+      cache = Jbuilder::Cache.new
+      cache.add("x", {}) do
+        cache.add("y", {}) do
+          cache.add("z", {}) do
+            { key: "value" }
+          end
+
+          ["y"]
+        end
+
+        "x"
+      end
+      cache.resolve
+    end
+
+    # cache miss:
+    #
+    #   1. x
+    #   2. y
+    #   3. z
+    #
+    # cache hit:
+    #
+    #   4. x + y + z
+    #   5. x + y + z
+    assert_equal 5, Rails.cache.fetch_multi_calls.length
+
+    # cache miss:
+    #
+    #   1. x
+    #   2. y
+    #   2. z
+    assert_equal 3, Rails.cache.write_calls.length
+  end
+
+  test "different options" do
+    3.times do
+      cache = Jbuilder::Cache.new
+      cache.add("x", {}) { "x" }
+      cache.add("y", {}) { [ "y" ] }
+      cache.add("z", { expires_in: 10.minutes }) { { key: "value" } }
+      cache.resolve
+    end
+
+    # cache miss:
+    #
+    #   1. x + y
+    #   2. z
+    #
+    # cache hit:
+    #
+    #   3. x + y
+    #   4. z
+    #   5. x + y
+    #   6. z
+    assert_equal 6, Rails.cache.fetch_multi_calls.length
+
+    # cache miss:
+    #
+    #   1. x
+    #   2. y
+    #   2. z
+    assert_equal 3, Rails.cache.write_calls.length
+  end
+end

test/jbuilder_template_test.rb

@@ -3,7 +3,6 @@
 require "active_model"
 require "action_view"
 require "action_view/testing/resolvers"
-require "active_support/cache"
 require "jbuilder/jbuilder_template"
 
 BLOG_POST_PARTIAL = <<-JBUILDER
@@ -50,12 +49,6 @@ def initialize(id, name)
   "_collection.json.jbuilder" => COLLECTION_PARTIAL
 }
 
-module Rails
-  def self.cache
-    @cache ||= ActiveSupport::Cache::MemoryStore.new
-  end
-end
-
 class JbuilderTemplateTest < ActionView::TestCase
   setup do
     @context = self
@@ -332,23 +325,6 @@ def assert_collection_rendered(result, context = nil)
     JBUILDER
   end
 
-  test "fragment caching instrumentation" do
-    undef_context_methods :fragment_name_with_digest, :cache_fragment_name
-
-    payloads = {}
-    ActiveSupport::Notifications.subscribe("read_fragment.action_controller") { |*args| payloads[:read_fragment] = args.last }
-    ActiveSupport::Notifications.subscribe("write_fragment.action_controller") { |*args| payloads[:write_fragment] = args.last }
-
-    jbuild <<-JBUILDER
-      json.cache! "cachekey" do
-        json.name "Cache"
-      end
-    JBUILDER
-
-    assert_equal "jbuilder/cachekey", payloads[:read_fragment][:key]
-    assert_equal "jbuilder/cachekey", payloads[:write_fragment][:key]
-  end
-
   test "current cache digest option accepts options" do
     undef_context_methods :fragment_name_with_digest
 

test/test_helper.rb

@@ -1,6 +1,7 @@
 require "bundler/setup"
 require "active_support"
 require 'active_support/core_ext/array/access'
+require "active_support/cache"
 require "rails/version"
 require "jbuilder"
 
@@ -10,7 +11,43 @@
   require "test/unit"
 end
 
-
 if ActiveSupport.respond_to?(:test_order=)
   ActiveSupport.test_order = :random
 end
+
+class MemoryStore < ActiveSupport::Cache::MemoryStore
+  attr_reader :fetch_multi_calls
+  attr_reader :write_calls
+
+  def initialize(*)
+    super
+
+    @fetch_multi_calls = []
+    @write_calls = []
+  end
+
+  def clear
+    @fetch_multi_calls.clear
+    @write_calls.clear
+
+    super
+  end
+
+  def fetch_multi(*args)
+    fetch_multi_calls << args
+
+    super
+  end
+
+  def write(*args)
+    write_calls << args
+
+    super
+  end
+end
+
+module Rails
+  def self.cache
+    @cache ||= MemoryStore.new
+  end
+end