From 7595790409740dd53fb016eca43f75c01bd6266f Mon Sep 17 00:00:00 2001
From: Jianbin Chen <jason.chen@shopify.com>
Date: Tue, 8 Oct 2024 22:44:44 -0700
Subject: [PATCH 1/7] Support deferred expiration of records and attributes for
 one-to-many associations

---
 Rakefile                                   | 24 ++++++++++++
 dev.yml                                    | 16 +++++---
 lib/identity_cache.rb                      | 20 ++++++++++
 lib/identity_cache/cache_fetcher.rb        |  9 +++++
 lib/identity_cache/cached/attribute.rb     | 18 +++++++--
 lib/identity_cache/cached/primary_index.rb |  6 ++-
 lib/identity_cache/memoized_cache_proxy.rb | 14 +++++++
 test/save_test.rb                          | 44 ++++++++++++++++++++++
 8 files changed, 141 insertions(+), 10 deletions(-)

diff --git a/Rakefile b/Rakefile
index 6212c524..1742ce5a 100644
--- a/Rakefile
+++ b/Rakefile
@@ -47,3 +47,27 @@ namespace :profile do
     ruby "./performance/profile.rb"
   end
 end
+
+namespace :db do
+  desc "Create the identity_cache_test database"
+  task :create do
+    require "mysql2"
+
+    config = {
+      host: ENV.fetch("MYSQL_HOST") || "localhost",
+      port: ENV.fetch("MYSQL_PORT") || 1037,
+      username: ENV.fetch("MYSQL_USER") || "root",
+      password: ENV.fetch("MYSQL_PASSWORD") || "",
+    }
+
+    begin
+      client = Mysql2::Client.new(config)
+      client.query("CREATE DATABASE IF NOT EXISTS identity_cache_test")
+      puts "Database 'identity_cache_test' created successfully. host: #{config[:host]}, port: #{config[:port]}"
+    rescue Mysql2::Error => e
+      puts "Error creating database: #{e.message}"
+    ensure
+      client&.close
+    end
+  end
+end
diff --git a/dev.yml b/dev.yml
index 58f0a341..98ec8ab3 100644
--- a/dev.yml
+++ b/dev.yml
@@ -5,6 +5,10 @@ up:
   - bundler
   - memcached
   - mysql
+  - custom:
+      name: create database identity_cache_test
+      met?: mysql -u root -h 127.0.0.1 -P 1037 -e "SHOW DATABASES;" | grep identity_cache_test
+      meet: bundle exec rake db:create
 
 commands:
   test:
@@ -12,7 +16,7 @@ commands:
       optional:
         argument: file
         optional: args...
-    desc: 'Run tests'
+    desc: "Run tests"
     run: |
       if [[ $# -eq 0 ]]; then
         bundle exec rake test
@@ -21,21 +25,21 @@ commands:
       fi
 
   style:
-    desc: 'Run rubocop checks'
+    desc: "Run rubocop checks"
     run: bundle exec rubocop "$@"
 
   check:
-    desc: 'Run tests and style checks'
+    desc: "Run tests and style checks"
     run: bundle exec rake test && bundle exec rubocop
 
   benchmark-cpu:
-    desc: 'Run the identity cache CPU benchmark'
+    desc: "Run the identity cache CPU benchmark"
     run: bundle exec rake benchmark:cpu
 
   profile:
-    desc: 'Profile IDC code'
+    desc: "Profile IDC code"
     run: bundle exec rake profile:run
 
   update-serialization-format:
-    desc: 'Update serialization format test fixture'
+    desc: "Update serialization format test fixture"
     run: bundle exec rake update_serialization_format
diff --git a/lib/identity_cache.rb b/lib/identity_cache.rb
index af492e56..026cc321 100644
--- a/lib/identity_cache.rb
+++ b/lib/identity_cache.rb
@@ -62,6 +62,8 @@ class InverseAssociationError < StandardError; end
 
   class NestedDeferredParentBlockError < StandardError; end
 
+  class NestedDeferredAttributeExpirationBlockError < StandardError; end
+
   class UnsupportedScopeError < StandardError; end
 
   class UnsupportedAssociationError < StandardError; end
@@ -240,6 +242,24 @@ def with_deferred_parent_expiration
       Thread.current[:idc_parent_records_for_cache_expiry].clear
     end
 
+    def with_deferred_attribute_expiration
+      raise NestedDeferredAttributeExpirationBlockError if Thread.current[:identity_cache_deferred_attribute_expiration]
+
+      Thread.current[:idc_deferred_attribute_expiration] = true
+      Thread.current[:idc_records_to_expire] = Set.new
+      Thread.current[:idc_attributes_to_expire] = Set.new
+
+      result = yield
+
+      Thread.current[:idc_deferred_attribute_expiration] = nil
+      IdentityCache.cache.delete_multi(Thread.current[:idc_records_to_expire])
+      IdentityCache.cache.delete_multi(Thread.current[:idc_attributes_to_expire])
+      result
+    ensure
+      Thread.current[:idc_deferred_attribute_expiration] = nil
+      Thread.current[:idc_attributes_to_expire].clear
+    end
+
     def with_fetch_read_only_records(value = true)
       old_value = Thread.current[:identity_cache_fetch_read_only_records]
       Thread.current[:identity_cache_fetch_read_only_records] = value
diff --git a/lib/identity_cache/cache_fetcher.rb b/lib/identity_cache/cache_fetcher.rb
index b08e2e69..e7f5c71f 100644
--- a/lib/identity_cache/cache_fetcher.rb
+++ b/lib/identity_cache/cache_fetcher.rb
@@ -60,6 +60,11 @@ def delete(key)
       @cache_backend.write(key, IdentityCache::DELETED, expires_in: IdentityCache::DELETED_TTL.seconds)
     end
 
+    def delete_multi(keys)
+      key_values = keys.map { |key| [key, IdentityCache::DELETED] }.to_h
+      @cache_backend.write_multi(key_values, expires_in: IdentityCache::DELETED_TTL.seconds)
+    end
+
     def clear
       @cache_backend.clear
     end
@@ -82,6 +87,10 @@ def fetch(key, fill_lock_duration: nil, lock_wait_tries: 2, &block)
       end
     end
 
+    def exist?(key)
+      @cache_backend.exist?(key)
+    end
+
     private
 
     def fetch_without_fill_lock(key)
diff --git a/lib/identity_cache/cached/attribute.rb b/lib/identity_cache/cached/attribute.rb
index c02dd9e8..9ffa0aa0 100644
--- a/lib/identity_cache/cached/attribute.rb
+++ b/lib/identity_cache/cached/attribute.rb
@@ -39,12 +39,24 @@ def expire(record)
 
         unless record.send(:was_new_record?)
           old_key = old_cache_key(record)
-          all_deleted = IdentityCache.cache.delete(old_key)
+
+          if Thread.current[:idc_deferred_attribute_expiration]
+            Thread.current[:idc_attributes_to_expire] << old_key
+            # defer the deletion, and don't block the following deletion
+            all_deleted = true
+          else
+            all_deleted = IdentityCache.cache.delete(old_key)
+          end
         end
         unless record.destroyed?
           new_key = new_cache_key(record)
           if new_key != old_key
-            all_deleted = IdentityCache.cache.delete(new_key) && all_deleted
+            if Thread.current[:idc_deferred_attribute_expiration]
+              Thread.current[:idc_attributes_to_expire] << new_key
+              all_deleted = true
+            else
+              all_deleted = IdentityCache.cache.delete(new_key) && all_deleted
+            end
           end
         end
 
@@ -152,9 +164,9 @@ def new_cache_key(record)
       end
 
       def old_cache_key(record)
+        changes = record.transaction_changed_attributes
         old_key_values = key_fields.map do |field|
           field_string = field.to_s
-          changes = record.transaction_changed_attributes
           if record.destroyed? && changes.key?(field_string)
             changes[field_string]
           elsif record.persisted? && changes.key?(field_string)
diff --git a/lib/identity_cache/cached/primary_index.rb b/lib/identity_cache/cached/primary_index.rb
index dff12dd3..538c95e0 100644
--- a/lib/identity_cache/cached/primary_index.rb
+++ b/lib/identity_cache/cached/primary_index.rb
@@ -45,7 +45,11 @@ def fetch_multi(ids)
 
       def expire(id)
         id = cast_id(id)
-        IdentityCache.cache.delete(cache_key(id))
+        if Thread.current[:idc_deferred_attribute_expiration]
+          Thread.current[:idc_records_to_expire] << cache_key(id)
+        else
+          IdentityCache.cache.delete(cache_key(id))
+        end
       end
 
       def cache_key(id)
diff --git a/lib/identity_cache/memoized_cache_proxy.rb b/lib/identity_cache/memoized_cache_proxy.rb
index 31e8c95a..da4a954c 100644
--- a/lib/identity_cache/memoized_cache_proxy.rb
+++ b/lib/identity_cache/memoized_cache_proxy.rb
@@ -70,6 +70,16 @@ def delete(key)
       end
     end
 
+    def delete_multi(keys)
+      memoizing = memoizing?
+      ActiveSupport::Notifications.instrument("cache_delete_multi.identity_cache", memoizing: memoizing) do
+        if memoizing
+          keys.each { |key| memoized_key_values.delete(key) }
+        end
+        @cache_fetcher.delete_multi(keys)
+      end
+    end
+
     def fetch(key, cache_fetcher_options = {}, &block)
       memo_misses = 0
       cache_misses = 0
@@ -137,6 +147,10 @@ def clear
       end
     end
 
+    def exist?(key)
+      @cache_fetcher.exist?(key)
+    end
+
     private
 
     EMPTY_ARRAY = [].freeze
diff --git a/test/save_test.rb b/test/save_test.rb
index 16b55ae7..dc954ee7 100644
--- a/test/save_test.rb
+++ b/test/save_test.rb
@@ -80,9 +80,53 @@ def test_expire_cache_works_in_a_transaction
     end
   end
 
+  def test_touch_with_separate_calls
+    @record1 = Item.create(title: "fooz", created_at: 1.second.ago, updated_at: 1.second.ago)
+    @record2 = Item.create(title: "barz", created_at: 1.second.ago, updated_at: 1.second.ago)
+    id_and_title_key1 = "\"#{@record1.id}\"/\"fooz\""
+    expect_cache_delete("#{NAMESPACE}attr:Item:id:id/title:#{cache_hash(id_and_title_key1)}")
+    expect_cache_delete("#{NAMESPACE}attr:Item:id:title:#{cache_hash('"fooz"')}")
+    id_and_title_key2 = "\"#{@record2.id}\"/\"barz\""
+    expect_cache_delete("#{NAMESPACE}attr:Item:id:id/title:#{cache_hash(id_and_title_key2)}")
+    expect_cache_delete("#{NAMESPACE}attr:Item:id:title:#{cache_hash('"barz"')}")
+    expect_cache_delete(@record1.primary_cache_index_key)
+    expect_cache_delete(@record2.primary_cache_index_key)
+
+    ActiveRecord::Base.transaction do
+      @record1.touch
+      @record2.touch
+    end
+  end
+
+  def test_touch_with_batched_calls
+    @record1 = Item.create(title: "fooz", created_at: 1.second.ago, updated_at: 1.second.ago)
+    @record2 = Item.create(title: "barz", created_at: 1.second.ago, updated_at: 1.second.ago)
+    id_and_title_key1 = "\"#{@record1.id}\"/\"fooz\""
+    id_and_title_key2 = "\"#{@record2.id}\"/\"barz\""
+    expect_cache_deletes([
+      "#{NAMESPACE}attr:Item:id:title:#{cache_hash('"fooz"')}",
+      "#{NAMESPACE}attr:Item:id:id/title:#{cache_hash(id_and_title_key1)}",
+      "#{NAMESPACE}attr:Item:id:title:#{cache_hash('"barz"')}",
+      "#{NAMESPACE}attr:Item:id:id/title:#{cache_hash(id_and_title_key2)}",
+    ])
+    expect_cache_deletes([@record1.primary_cache_index_key, @record2.primary_cache_index_key])
+
+    IdentityCache.with_deferred_attribute_expiration do
+      ActiveRecord::Base.transaction do
+        @record1.touch
+        @record2.touch
+      end
+    end
+  end
+
   private
 
   def expect_cache_delete(key)
     @backend.expects(:write).with(key, IdentityCache::DELETED, anything)
   end
+
+  def expect_cache_deletes(keys)
+    key_values = keys.map { |key| [key, IdentityCache::DELETED] }
+    @backend.expects(:write_multi).with(key_values, anything)
+  end
 end

From 9f4bed3be41ada2feb42f426e187720a7961d8c9 Mon Sep 17 00:00:00 2001
From: Jianbin Chen <jason.chen@shopify.com>
Date: Wed, 9 Oct 2024 14:35:45 -0700
Subject: [PATCH 2/7] Merge the two methods into one

---
 lib/identity_cache.rb                         | 56 +++++++------------
 lib/identity_cache/cached/attribute.rb        |  4 +-
 lib/identity_cache/cached/primary_index.rb    |  4 +-
 lib/identity_cache/parent_model_expiration.rb |  4 +-
 test/index_cache_test.rb                      | 44 +++++++--------
 test/save_test.rb                             | 12 +---
 test/test_helper.rb                           | 10 ++++
 7 files changed, 58 insertions(+), 76 deletions(-)

diff --git a/lib/identity_cache.rb b/lib/identity_cache.rb
index 026cc321..78e750bb 100644
--- a/lib/identity_cache.rb
+++ b/lib/identity_cache.rb
@@ -60,9 +60,7 @@ class AssociationError < StandardError; end
 
   class InverseAssociationError < StandardError; end
 
-  class NestedDeferredParentBlockError < StandardError; end
-
-  class NestedDeferredAttributeExpirationBlockError < StandardError; end
+  class NestedDeferredCacheExpirationBlockError < StandardError; end
 
   class UnsupportedScopeError < StandardError; end
 
@@ -204,59 +202,45 @@ def fetch_multi(*keys)
       result
     end
 
-    # Executes a block with deferred parent expiration, ensuring that the parent
-    # records' cache expiration is deferred until the block completes. When the block
-    # completes, it triggers expiration of the primary index for the parent records.
-    # Raises a NestedDeferredParentBlockError if a deferred parent expiration block
-    # is already active on the current thread.
+    # Executes a block with deferred cache expiration, ensuring that the records' (parent,
+    # children and attributes) cache expiration is deferred until the block completes. When
+    # the block completes, it issues delete_multi calls for all the records and attributes
+    # that were marked for expiration.
     #
     # == Parameters:
     # No parameters.
     #
     # == Raises:
-    # NestedDeferredParentBlockError if a deferred parent expiration block is already active.
+    # NestedDeferredCacheExpirationBlockError if a deferred cache expiration block is already active.
     #
     # == Yield:
-    # Runs the provided block with deferred parent expiration.
+    # Runs the provided block with deferred cache expiration.
     #
     # == Returns:
     # The result of executing the provided block.
     #
     # == Ensures:
-    # Cleans up thread-local variables related to deferred parent expiration regardless
+    # Cleans up thread-local variables related to deferred cache expiration regardless
     # of whether the block raises an exception.
-    def with_deferred_parent_expiration
-      raise NestedDeferredParentBlockError if Thread.current[:idc_deferred_parent_expiration]
-
-      Thread.current[:idc_deferred_parent_expiration] = true
-      Thread.current[:idc_parent_records_for_cache_expiry] = Set.new
-
-      result = yield
-
-      Thread.current[:idc_deferred_parent_expiration] = nil
-      Thread.current[:idc_parent_records_for_cache_expiry].each(&:expire_primary_index)
-
-      result
-    ensure
-      Thread.current[:idc_deferred_parent_expiration] = nil
-      Thread.current[:idc_parent_records_for_cache_expiry].clear
-    end
-
-    def with_deferred_attribute_expiration
-      raise NestedDeferredAttributeExpirationBlockError if Thread.current[:identity_cache_deferred_attribute_expiration]
+    def with_deferred_expiration
+      raise NestedDeferredCacheExpirationBlockError if Thread.current[:idc_deferred_expiration]
 
-      Thread.current[:idc_deferred_attribute_expiration] = true
-      Thread.current[:idc_records_to_expire] = Set.new
+      Thread.current[:idc_deferred_expiration] = true
+      Thread.current[:idc_parent_records_to_expire] = Set.new
+      Thread.current[:idc_child_records_to_expire] = Set.new
       Thread.current[:idc_attributes_to_expire] = Set.new
 
       result = yield
 
-      Thread.current[:idc_deferred_attribute_expiration] = nil
-      IdentityCache.cache.delete_multi(Thread.current[:idc_records_to_expire])
-      IdentityCache.cache.delete_multi(Thread.current[:idc_attributes_to_expire])
+      Thread.current[:idc_deferred_expiration] = nil
+      IdentityCache.cache.delete_multi(Thread.current[:idc_parent_records_to_expire]) if Thread.current[:idc_parent_records_to_expire].any?
+      IdentityCache.cache.delete_multi(Thread.current[:idc_child_records_to_expire]) if Thread.current[:idc_child_records_to_expire].any?
+      IdentityCache.cache.delete_multi(Thread.current[:idc_attributes_to_expire]) if Thread.current[:idc_attributes_to_expire].any?
       result
     ensure
-      Thread.current[:idc_deferred_attribute_expiration] = nil
+      Thread.current[:idc_deferred_expiration] = nil
+      Thread.current[:idc_parent_records_to_expire].clear
+      Thread.current[:idc_child_records_to_expire].clear
       Thread.current[:idc_attributes_to_expire].clear
     end
 
diff --git a/lib/identity_cache/cached/attribute.rb b/lib/identity_cache/cached/attribute.rb
index 9ffa0aa0..bbaec914 100644
--- a/lib/identity_cache/cached/attribute.rb
+++ b/lib/identity_cache/cached/attribute.rb
@@ -40,7 +40,7 @@ def expire(record)
         unless record.send(:was_new_record?)
           old_key = old_cache_key(record)
 
-          if Thread.current[:idc_deferred_attribute_expiration]
+          if Thread.current[:idc_deferred_expiration]
             Thread.current[:idc_attributes_to_expire] << old_key
             # defer the deletion, and don't block the following deletion
             all_deleted = true
@@ -51,7 +51,7 @@ def expire(record)
         unless record.destroyed?
           new_key = new_cache_key(record)
           if new_key != old_key
-            if Thread.current[:idc_deferred_attribute_expiration]
+            if Thread.current[:idc_deferred_expiration]
               Thread.current[:idc_attributes_to_expire] << new_key
               all_deleted = true
             else
diff --git a/lib/identity_cache/cached/primary_index.rb b/lib/identity_cache/cached/primary_index.rb
index 538c95e0..3abd04a5 100644
--- a/lib/identity_cache/cached/primary_index.rb
+++ b/lib/identity_cache/cached/primary_index.rb
@@ -45,8 +45,8 @@ def fetch_multi(ids)
 
       def expire(id)
         id = cast_id(id)
-        if Thread.current[:idc_deferred_attribute_expiration]
-          Thread.current[:idc_records_to_expire] << cache_key(id)
+        if Thread.current[:idc_deferred_expiration]
+          Thread.current[:idc_child_records_to_expire] << cache_key(id)
         else
           IdentityCache.cache.delete(cache_key(id))
         end
diff --git a/lib/identity_cache/parent_model_expiration.rb b/lib/identity_cache/parent_model_expiration.rb
index 499e6d0b..a26fff44 100644
--- a/lib/identity_cache/parent_model_expiration.rb
+++ b/lib/identity_cache/parent_model_expiration.rb
@@ -47,8 +47,8 @@ def expire_parent_caches
       add_parents_to_cache_expiry_set(parents_to_expire)
       parents_to_expire.select! { |parent| parent.class.primary_cache_index_enabled }
       parents_to_expire.reduce(true) do |all_expired, parent|
-        if Thread.current[:idc_deferred_parent_expiration]
-          Thread.current[:idc_parent_records_for_cache_expiry] << parent
+        if Thread.current[:idc_deferred_expiration]
+          Thread.current[:idc_parent_records_to_expire] << parent.cache_key
           next parent
         end
         parent.expire_primary_index && all_expired
diff --git a/test/index_cache_test.rb b/test/index_cache_test.rb
index bb3a48ea..1fbf3f17 100644
--- a/test/index_cache_test.rb
+++ b/test/index_cache_test.rb
@@ -166,38 +166,36 @@ def test_unique_cache_index_with_non_id_primary_key
     assert_equal(123, KeyedRecord.fetch_by_value("a").id)
   end
 
-  def test_with_deferred_parent_expiration_expires_parent_index_once
+  def test_with_deferred_expiration_for_parent_records_expires_parent_index_once
     Item.send(:cache_has_many, :associated_records, embed: true)
 
     @parent = Item.create!(title: "bob")
     @records = @parent.associated_records.create!([{ name: "foo" }, { name: "bar" }, { name: "baz" }])
 
-    @memcached_spy = Spy.on(backend, :write).and_call_through
-
+    @memcached_spy_write_multi = Spy.on(backend, :write_multi).and_call_through
     expected_item_expiration_count = Array(@parent).count
     expected_associated_record_expiration_count = @records.count
 
     expected_return_value = "Some text that we expect to see returned from the block"
 
-    result = IdentityCache.with_deferred_parent_expiration do
+    result = IdentityCache.with_deferred_expiration do
       @parent.transaction do
         @parent.associated_records.destroy_all
       end
-      assert_equal(expected_associated_record_expiration_count, @memcached_spy.calls.count)
       expected_return_value
     end
 
-    expired_cache_keys = @memcached_spy.calls.map(&:args).map(&:first)
-    item_expiration_count = expired_cache_keys.count { _1.include?("Item") }
-    associated_record_expiration_count = expired_cache_keys.count { _1.include?("AssociatedRecord") }
+    all_keys = @memcached_spy_write_multi.calls.flat_map { |call| call.args.first.keys }
+    item_expiration_count = all_keys.count { _1.start_with?("items/") }
+    associated_record_expiration_count = all_keys.count { _1.include?(":AssociatedRecord:") }
 
-    assert_operator(@memcached_spy.calls.count, :>, 0)
+    assert_operator(@memcached_spy_write_multi.calls.count, :>, 0)
     assert_equal(expected_item_expiration_count, item_expiration_count)
     assert_equal(expected_associated_record_expiration_count, associated_record_expiration_count)
     assert_equal(expected_return_value, result)
   end
 
-  def test_double_nested_deferred_parent_expiration_will_raise_error
+  def test_double_nested_deferred_expiration_for_parent_records_will_raise_error
     Item.send(:cache_has_many, :associated_records, embed: true)
 
     @parent = Item.create!(title: "bob")
@@ -205,9 +203,9 @@ def test_double_nested_deferred_parent_expiration_will_raise_error
 
     @memcached_spy = Spy.on(backend, :write).and_call_through
 
-    assert_raises(IdentityCache::NestedDeferredParentBlockError) do
-      IdentityCache.with_deferred_parent_expiration do
-        IdentityCache.with_deferred_parent_expiration do
+    assert_raises(IdentityCache::NestedDeferredCacheExpirationBlockError) do
+      IdentityCache.with_deferred_expiration do
+        IdentityCache.with_deferred_expiration do
           @parent.transaction do
             @parent.associated_records.destroy_all
           end
@@ -218,7 +216,7 @@ def test_double_nested_deferred_parent_expiration_will_raise_error
     assert_equal(0, @memcached_spy.calls.count)
   end
 
-  def test_deep_association_with_deferred_parent_expiration_expires_parent_once
+  def test_deep_association_with_deferred_expiration_expires_parent_once
     AssociatedRecord.send(:has_many, :deeply_associated_records, dependent: :destroy)
     Item.send(:cache_has_many, :associated_records, embed: true)
 
@@ -232,27 +230,27 @@ def test_deep_association_with_deferred_parent_expiration_expires_parent_once
       ])
     end
 
-    @memcached_spy = Spy.on(backend, :write).and_call_through
+    @memcached_spy_write_multi = Spy.on(backend, :write_multi).and_call_through
 
     expected_item_expiration_count = Array(@parent).count
     expected_associated_record_expiration_count = @records.count
     expected_deeply_associated_record_expiration_count = @records.flat_map(&:deeply_associated_records).count
 
-    IdentityCache.with_deferred_parent_expiration do
+    IdentityCache.with_deferred_expiration do
       @parent.transaction do
         @parent.associated_records.destroy_all
       end
     end
 
-    expired_cache_keys = @memcached_spy.calls.map(&:args).map(&:first)
-    item_expiration_count = expired_cache_keys.count { _1.include?("Item") }
-    associated_record_expiration_count = expired_cache_keys.count { _1.include?(":AssociatedRecord:") }
-    deeply_associated_record_expiration_count = expired_cache_keys.count { _1.include?("DeeplyAssociatedRecord") }
+    all_keys = @memcached_spy_write_multi.calls.flat_map { |call| call.args.first.keys }
+    item_expiration_count = all_keys.count { |key| key.start_with?("items/") }
+    associated_record_keys = all_keys.count { |key| key.include?(":AssociatedRecord:") }
+    deeply_associated_record_keys = all_keys.count { |key| key.include?(":DeeplyAssociatedRecord:") }
 
-    assert_operator(@memcached_spy.calls.count, :>, 0)
+    assert_operator(@memcached_spy_write_multi.calls.count, :>, 0)
     assert_equal(expected_item_expiration_count, item_expiration_count)
-    assert_equal(expected_associated_record_expiration_count, associated_record_expiration_count)
-    assert_equal(expected_deeply_associated_record_expiration_count, deeply_associated_record_expiration_count)
+    assert_equal(expected_associated_record_expiration_count, associated_record_keys)
+    assert_equal(expected_deeply_associated_record_expiration_count, deeply_associated_record_keys)
   end
 
   private
diff --git a/test/save_test.rb b/test/save_test.rb
index dc954ee7..2aa16253 100644
--- a/test/save_test.rb
+++ b/test/save_test.rb
@@ -111,7 +111,7 @@ def test_touch_with_batched_calls
     ])
     expect_cache_deletes([@record1.primary_cache_index_key, @record2.primary_cache_index_key])
 
-    IdentityCache.with_deferred_attribute_expiration do
+    IdentityCache.with_deferred_expiration do
       ActiveRecord::Base.transaction do
         @record1.touch
         @record2.touch
@@ -119,14 +119,4 @@ def test_touch_with_batched_calls
     end
   end
 
-  private
-
-  def expect_cache_delete(key)
-    @backend.expects(:write).with(key, IdentityCache::DELETED, anything)
-  end
-
-  def expect_cache_deletes(keys)
-    key_values = keys.map { |key| [key, IdentityCache::DELETED] }
-    @backend.expects(:write_multi).with(key_values, anything)
-  end
 end
diff --git a/test/test_helper.rb b/test/test_helper.rb
index ed5df3b8..a37c7947 100644
--- a/test/test_helper.rb
+++ b/test/test_helper.rb
@@ -9,6 +9,7 @@
 require "helpers/active_record_objects"
 require "helpers/mocked_cache_backend"
 require "spy/integration"
+require "byebug"
 
 require File.dirname(__FILE__) + "/../lib/identity_cache"
 
@@ -128,6 +129,15 @@ def assert_memcache_operations(num, &block)
       ret
     end
 
+    def expect_cache_delete(key)
+      @backend.expects(:write).with(key, IdentityCache::DELETED, anything)
+    end
+
+    def expect_cache_deletes(keys)
+      key_values = keys.map { |key| [key, IdentityCache::DELETED] }.to_h
+      @backend.expects(:write_multi).with(key_values, anything)
+    end
+
     def assert_no_queries(**subscribe_opts, &block)
       subscribe_to_sql_queries(->(sql) { raise "Unexpected SQL query: #{sql}" }, **subscribe_opts, &block)
     end

From 68f729696002c85fda76ae9f78d9d20ce05cac20 Mon Sep 17 00:00:00 2001
From: Jianbin Chen <jason.chen@shopify.com>
Date: Wed, 9 Oct 2024 14:38:42 -0700
Subject: [PATCH 3/7] Remove byedebug

---
 test/test_helper.rb | 1 -
 1 file changed, 1 deletion(-)

diff --git a/test/test_helper.rb b/test/test_helper.rb
index a37c7947..5b00c11c 100644
--- a/test/test_helper.rb
+++ b/test/test_helper.rb
@@ -9,7 +9,6 @@
 require "helpers/active_record_objects"
 require "helpers/mocked_cache_backend"
 require "spy/integration"
-require "byebug"
 
 require File.dirname(__FILE__) + "/../lib/identity_cache"
 

From 6954b79b438d30e0a625ba77427c2fae89a2ceea Mon Sep 17 00:00:00 2001
From: Jianbin Chen <jason.chen@shopify.com>
Date: Wed, 9 Oct 2024 14:46:06 -0700
Subject: [PATCH 4/7] Fix rubocop

---
 lib/identity_cache.rb | 20 +++++++++++++++++---
 test/save_test.rb     |  1 -
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/lib/identity_cache.rb b/lib/identity_cache.rb
index 78e750bb..fcead9ef 100644
--- a/lib/identity_cache.rb
+++ b/lib/identity_cache.rb
@@ -233,9 +233,23 @@ def with_deferred_expiration
       result = yield
 
       Thread.current[:idc_deferred_expiration] = nil
-      IdentityCache.cache.delete_multi(Thread.current[:idc_parent_records_to_expire]) if Thread.current[:idc_parent_records_to_expire].any?
-      IdentityCache.cache.delete_multi(Thread.current[:idc_child_records_to_expire]) if Thread.current[:idc_child_records_to_expire].any?
-      IdentityCache.cache.delete_multi(Thread.current[:idc_attributes_to_expire]) if Thread.current[:idc_attributes_to_expire].any?
+      if Thread.current[:idc_parent_records_to_expire].any?
+        IdentityCache.cache.delete_multi(
+          Thread.current[:idc_parent_records_to_expire]
+        )
+      end
+
+      if Thread.current[:idc_child_records_to_expire].any?
+        IdentityCache.cache.delete_multi(
+          Thread.current[:idc_child_records_to_expire]
+        )
+      end
+
+      if Thread.current[:idc_attributes_to_expire].any?
+        IdentityCache.cache.delete_multi(
+          Thread.current[:idc_attributes_to_expire]
+        )
+      end
       result
     ensure
       Thread.current[:idc_deferred_expiration] = nil
diff --git a/test/save_test.rb b/test/save_test.rb
index 2aa16253..976f8216 100644
--- a/test/save_test.rb
+++ b/test/save_test.rb
@@ -118,5 +118,4 @@ def test_touch_with_batched_calls
       end
     end
   end
-
 end

From 2f189acc021269bb3b8a06ad9df38575242a9ecf Mon Sep 17 00:00:00 2001
From: Jianbin Chen <jason.chen@shopify.com>
Date: Wed, 9 Oct 2024 15:20:36 -0700
Subject: [PATCH 5/7] Merge the parent and children into one set

---
 lib/identity_cache.rb                         | 17 ++++-------------
 lib/identity_cache/cached/primary_index.rb    |  2 +-
 lib/identity_cache/parent_model_expiration.rb |  4 ----
 test/index_cache_test.rb                      | 16 +++++++++-------
 4 files changed, 14 insertions(+), 25 deletions(-)

diff --git a/lib/identity_cache.rb b/lib/identity_cache.rb
index fcead9ef..7cd3e142 100644
--- a/lib/identity_cache.rb
+++ b/lib/identity_cache.rb
@@ -226,25 +226,17 @@ def with_deferred_expiration
       raise NestedDeferredCacheExpirationBlockError if Thread.current[:idc_deferred_expiration]
 
       Thread.current[:idc_deferred_expiration] = true
-      Thread.current[:idc_parent_records_to_expire] = Set.new
-      Thread.current[:idc_child_records_to_expire] = Set.new
+      Thread.current[:idc_records_to_expire] = Set.new
       Thread.current[:idc_attributes_to_expire] = Set.new
 
       result = yield
 
       Thread.current[:idc_deferred_expiration] = nil
-      if Thread.current[:idc_parent_records_to_expire].any?
+      if Thread.current[:idc_records_to_expire].any?
         IdentityCache.cache.delete_multi(
-          Thread.current[:idc_parent_records_to_expire]
+          Thread.current[:idc_records_to_expire]
         )
       end
-
-      if Thread.current[:idc_child_records_to_expire].any?
-        IdentityCache.cache.delete_multi(
-          Thread.current[:idc_child_records_to_expire]
-        )
-      end
-
       if Thread.current[:idc_attributes_to_expire].any?
         IdentityCache.cache.delete_multi(
           Thread.current[:idc_attributes_to_expire]
@@ -253,8 +245,7 @@ def with_deferred_expiration
       result
     ensure
       Thread.current[:idc_deferred_expiration] = nil
-      Thread.current[:idc_parent_records_to_expire].clear
-      Thread.current[:idc_child_records_to_expire].clear
+      Thread.current[:idc_records_to_expire].clear
       Thread.current[:idc_attributes_to_expire].clear
     end
 
diff --git a/lib/identity_cache/cached/primary_index.rb b/lib/identity_cache/cached/primary_index.rb
index 3abd04a5..fcec3146 100644
--- a/lib/identity_cache/cached/primary_index.rb
+++ b/lib/identity_cache/cached/primary_index.rb
@@ -46,7 +46,7 @@ def fetch_multi(ids)
       def expire(id)
         id = cast_id(id)
         if Thread.current[:idc_deferred_expiration]
-          Thread.current[:idc_child_records_to_expire] << cache_key(id)
+          Thread.current[:idc_records_to_expire] << cache_key(id)
         else
           IdentityCache.cache.delete(cache_key(id))
         end
diff --git a/lib/identity_cache/parent_model_expiration.rb b/lib/identity_cache/parent_model_expiration.rb
index a26fff44..1993a21c 100644
--- a/lib/identity_cache/parent_model_expiration.rb
+++ b/lib/identity_cache/parent_model_expiration.rb
@@ -47,10 +47,6 @@ def expire_parent_caches
       add_parents_to_cache_expiry_set(parents_to_expire)
       parents_to_expire.select! { |parent| parent.class.primary_cache_index_enabled }
       parents_to_expire.reduce(true) do |all_expired, parent|
-        if Thread.current[:idc_deferred_expiration]
-          Thread.current[:idc_parent_records_to_expire] << parent.cache_key
-          next parent
-        end
         parent.expire_primary_index && all_expired
       end
     end
diff --git a/test/index_cache_test.rb b/test/index_cache_test.rb
index 1fbf3f17..e9ffaa1f 100644
--- a/test/index_cache_test.rb
+++ b/test/index_cache_test.rb
@@ -186,10 +186,10 @@ def test_with_deferred_expiration_for_parent_records_expires_parent_index_once
     end
 
     all_keys = @memcached_spy_write_multi.calls.flat_map { |call| call.args.first.keys }
-    item_expiration_count = all_keys.count { _1.start_with?("items/") }
-    associated_record_expiration_count = all_keys.count { _1.include?(":AssociatedRecord:") }
+    item_expiration_count = all_keys.count { _1.include?(":blob:Item:") }
+    associated_record_expiration_count = all_keys.count { _1.include?(":blob:AssociatedRecord:") }
 
-    assert_operator(@memcached_spy_write_multi.calls.count, :>, 0)
+    assert_equal(1, @memcached_spy_write_multi.calls.count)
     assert_equal(expected_item_expiration_count, item_expiration_count)
     assert_equal(expected_associated_record_expiration_count, associated_record_expiration_count)
     assert_equal(expected_return_value, result)
@@ -202,6 +202,7 @@ def test_double_nested_deferred_expiration_for_parent_records_will_raise_error
     @records = @parent.associated_records.create!([{ name: "foo" }, { name: "bar" }, { name: "baz" }])
 
     @memcached_spy = Spy.on(backend, :write).and_call_through
+    @memcached_spy_write_multi = Spy.on(backend, :write_multi).and_call_through
 
     assert_raises(IdentityCache::NestedDeferredCacheExpirationBlockError) do
       IdentityCache.with_deferred_expiration do
@@ -214,6 +215,7 @@ def test_double_nested_deferred_expiration_for_parent_records_will_raise_error
     end
 
     assert_equal(0, @memcached_spy.calls.count)
+    assert_equal(0, @memcached_spy_write_multi.calls.count)
   end
 
   def test_deep_association_with_deferred_expiration_expires_parent_once
@@ -243,11 +245,11 @@ def test_deep_association_with_deferred_expiration_expires_parent_once
     end
 
     all_keys = @memcached_spy_write_multi.calls.flat_map { |call| call.args.first.keys }
-    item_expiration_count = all_keys.count { |key| key.start_with?("items/") }
-    associated_record_keys = all_keys.count { |key| key.include?(":AssociatedRecord:") }
-    deeply_associated_record_keys = all_keys.count { |key| key.include?(":DeeplyAssociatedRecord:") }
+    item_expiration_count = all_keys.count { |key| key.include?(":blob:Item:") }
+    associated_record_keys = all_keys.count { |key| key.include?(":blob:AssociatedRecord:") }
+    deeply_associated_record_keys = all_keys.count { |key| key.include?(":blob:DeeplyAssociatedRecord:") }
 
-    assert_operator(@memcached_spy_write_multi.calls.count, :>, 0)
+    assert_equal(1, @memcached_spy_write_multi.calls.count)
     assert_equal(expected_item_expiration_count, item_expiration_count)
     assert_equal(expected_associated_record_expiration_count, associated_record_keys)
     assert_equal(expected_deeply_associated_record_expiration_count, deeply_associated_record_keys)

From afac6a74fe3a794cf6eed9609c847a13d0994cc4 Mon Sep 17 00:00:00 2001
From: Jianbin Chen <jason.chen@shopify.com>
Date: Wed, 9 Oct 2024 15:23:44 -0700
Subject: [PATCH 6/7] Remove exist

---
 lib/identity_cache/cache_fetcher.rb        | 4 ----
 lib/identity_cache/memoized_cache_proxy.rb | 4 ----
 2 files changed, 8 deletions(-)

diff --git a/lib/identity_cache/cache_fetcher.rb b/lib/identity_cache/cache_fetcher.rb
index e7f5c71f..3573a628 100644
--- a/lib/identity_cache/cache_fetcher.rb
+++ b/lib/identity_cache/cache_fetcher.rb
@@ -87,10 +87,6 @@ def fetch(key, fill_lock_duration: nil, lock_wait_tries: 2, &block)
       end
     end
 
-    def exist?(key)
-      @cache_backend.exist?(key)
-    end
-
     private
 
     def fetch_without_fill_lock(key)
diff --git a/lib/identity_cache/memoized_cache_proxy.rb b/lib/identity_cache/memoized_cache_proxy.rb
index da4a954c..d967a149 100644
--- a/lib/identity_cache/memoized_cache_proxy.rb
+++ b/lib/identity_cache/memoized_cache_proxy.rb
@@ -147,10 +147,6 @@ def clear
       end
     end
 
-    def exist?(key)
-      @cache_fetcher.exist?(key)
-    end
-
     private
 
     EMPTY_ARRAY = [].freeze

From 8c964aeb5dc3ef72255d05f53042ccec399a5c72 Mon Sep 17 00:00:00 2001
From: Jianbin Chen <jason.chen@shopify.com>
Date: Thu, 10 Oct 2024 08:14:43 -0700
Subject: [PATCH 7/7] Add change log and bump the version

---
 CHANGELOG.md                  | 4 ++++
 Gemfile.lock                  | 2 +-
 lib/identity_cache/version.rb | 2 +-
 3 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2d094708..bdd12c7c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,10 @@
 
 ## Unreleased
 
+## 1.6.2
+
+- Support deferred expiry of associations and attributes. Add a rake task to create test database.
+
 ## 1.6.1
 
 - Fix deprecation warnings on Active Record 7.2. (#575)
diff --git a/Gemfile.lock b/Gemfile.lock
index 8e4f693d..6b392fac 100644
--- a/Gemfile.lock
+++ b/Gemfile.lock
@@ -8,7 +8,7 @@ GIT
 PATH
   remote: .
   specs:
-    identity_cache (1.6.1)
+    identity_cache (1.6.2)
       activerecord (>= 7.0)
       ar_transaction_changes (~> 1.1)
 
diff --git a/lib/identity_cache/version.rb b/lib/identity_cache/version.rb
index 3f4ffb37..f100f710 100644
--- a/lib/identity_cache/version.rb
+++ b/lib/identity_cache/version.rb
@@ -1,6 +1,6 @@
 # frozen_string_literal: true
 
 module IdentityCache
-  VERSION = "1.6.1"
+  VERSION = "1.6.2"
   CACHE_VERSION = 8
 end