From d552b976affa1641259f8b15037f26df0df3c9ec Mon Sep 17 00:00:00 2001
From: Gabriel Sobrinho <gabriel.sobrinho@gmail.com>
Date: Sat, 11 Nov 2017 17:23:35 -0200
Subject: [PATCH] Use fetch_multi for multiple cache blocks

This will improve the performance since we are going to hit the backend
just once to read all keys, as long the cache adapter implements fetch
multi support like dalli.

For example:

    json.cache! :x do
      json.x true
    end

    json.cache! :y do
      json.y true
    end

    json.cache! :z do
      json.z true
    end

This example was hitting the memcached 6 times on cache miss:

  1. read x
  2. write x
  3. read y
  4. write y
  5. read z
  6. write z

And 3 times on cache hit:

  1. read x
  2. read y
  3. read z

After this change, 4 times on cache miss:

  1. read multi x,y,z
  2. write x
  3. write y
  4. write z

And 1 time on cache hit:

  1. read multi x,y,z

Note that in the case of different options, one read multi will be made
per each options, i.e.:

    json.cache! :x do
      json.x true
    end

    json.cache! :y do
      json.y true
    end

    json.cache! :z, expires_in: 10.minutes do
      json.z true
    end

    json.cache! :w, expires_in: 10.minutes do
      json.w true
    end

In the case of cache miss:

  1. read multi x,y
  2. write x
  3. write y
  4. read multi z,w
  5. write z
  5. write w

In the case of cache hit:

  1. read multi x,y
  2. read multi z,w

That's because Rails.cache.fetch_multi signature is limited to use the
same options for all given keys.

And for last, nested cache calls are allowed and will follow recursively
to accomplish the same behavior, i.e.:

  json.cache! :x do
    json.x true

    json.cache! :y do
      json.y true
    end

    json.cache! :z do
      json.z true
    end
  end

  json.cache! :w do
    json.w true
  end

In the case of cache miss:

  1. read multi x,w
  2. read multi y,z
  3. write y
  4. write z
  5. write x
  6. write w

In the case of cache hit:

  1. read multi x,w

The same rule of options will be applied, if you have different options,
one hit per options.

This is the result of an investigation in application that was spending
15% of the time by hitting the memcached multiple times.

We were able to reduce the memcached time to 1% of the request by using
this algorithm.

Thanks to @samflores for helping me on the initial idea.
---
 lib/jbuilder.rb                   |   1 +
 lib/jbuilder/cache.rb             |  37 +++++++++
 lib/jbuilder/jbuilder_template.rb |  32 +++-----
 test/jbuilder_cache_test.rb       | 132 ++++++++++++++++++++++++++++++
 test/jbuilder_template_test.rb    |  24 ------
 test/test_helper.rb               |  39 ++++++++-
 6 files changed, 220 insertions(+), 45 deletions(-)
 create mode 100644 lib/jbuilder/cache.rb
 create mode 100644 test/jbuilder_cache_test.rb

diff --git a/lib/jbuilder.rb b/lib/jbuilder.rb
index d84eab2..18c50db 100644
--- a/lib/jbuilder.rb
+++ b/lib/jbuilder.rb
@@ -1,4 +1,5 @@
 require 'jbuilder/jbuilder'
+require 'jbuilder/cache'
 require 'jbuilder/blank'
 require 'jbuilder/key_formatter'
 require 'jbuilder/errors'
diff --git a/lib/jbuilder/cache.rb b/lib/jbuilder/cache.rb
new file mode 100644
index 0000000..c48f9f1
--- /dev/null
+++ b/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
diff --git a/lib/jbuilder/jbuilder_template.rb b/lib/jbuilder/jbuilder_template.rb
index 09a67c3..7ee4cb8 100644
--- a/lib/jbuilder/jbuilder_template.rb
+++ b/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)
diff --git a/test/jbuilder_cache_test.rb b/test/jbuilder_cache_test.rb
new file mode 100644
index 0000000..40c2203
--- /dev/null
+++ b/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
diff --git a/test/jbuilder_template_test.rb b/test/jbuilder_template_test.rb
index d8bfc41..c8c0288 100644
--- a/test/jbuilder_template_test.rb
+++ b/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
 
diff --git a/test/test_helper.rb b/test/test_helper.rb
index 330ce31..b27c342 100644
--- a/test/test_helper.rb
+++ b/test/test_helper.rb
@@ -1,5 +1,6 @@
 require "bundler/setup"
 require "active_support"
+require "active_support/cache"
 require "rails/version"
 
 if Rails::VERSION::STRING > "4.0"
@@ -8,7 +9,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