From 4a07e6ab2c9eaa662baffc076b6bfd91bb2a3bae Mon Sep 17 00:00:00 2001
From: Efstathios Stivaros <stath.stivaros@shopify.com>
Date: Mon, 10 Jun 2024 16:23:02 +0300
Subject: [PATCH] Add mechanism to prevent nested blocks

This could have caused the cache to have been overwritten
---
 lib/identity_cache.rb    |  3 +++
 test/index_cache_test.rb | 28 ++++++++++++++++++++++++++++
 2 files changed, 31 insertions(+)

diff --git a/lib/identity_cache.rb b/lib/identity_cache.rb
index 7a3ce7fd..91d43bc8 100644
--- a/lib/identity_cache.rb
+++ b/lib/identity_cache.rb
@@ -60,6 +60,8 @@ class AssociationError < StandardError; end
 
   class InverseAssociationError < StandardError; end
 
+  class NestedDeferredParentBlockError < StandardError; end
+
   class UnsupportedScopeError < StandardError; end
 
   class UnsupportedAssociationError < StandardError; end
@@ -199,6 +201,7 @@ def with_deferred_parent_expiration
 
       yield
 
+      Thread.current[:deferred_parent_expiration] = nil
       Thread.current[:parent_records_for_cache_expiry].each(&:expire_primary_index)
     end
 
diff --git a/test/index_cache_test.rb b/test/index_cache_test.rb
index bede53c4..c4afede7 100644
--- a/test/index_cache_test.rb
+++ b/test/index_cache_test.rb
@@ -191,6 +191,32 @@ def test_with_deferred_parent_expiration
     assert_operator(@memcached_spy.calls.count, :>, 0)
     assert_equal(expected_item_expiration_count, item_expiration_count)
     assert_equal(expected_associated_record_expiration_count, associated_record_expiration_count)
+  ensure
+    Thread.current[:deferred_parent_expiration] = nil
+  end
+
+  def test_double_nested_deferred_parent_expiration
+    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
+
+    assert_raises(IdentityCache::NestedDeferredParentBlockError) do
+      IdentityCache.with_deferred_parent_expiration do
+        IdentityCache.with_deferred_parent_expiration do
+          @parent.transaction do
+            @parent.associated_records.destroy_all
+          end
+          assert_equal(expected_associated_record_expiration_count, @memcached_spy.calls.count)
+        end
+      end
+    end
+
+    assert_equal(0, @memcached_spy.calls.count)
+  ensure
+    Thread.current[:deferred_parent_expiration] = nil
   end
 
   def test_deep_association_with_deferred_parent_expiration
@@ -228,6 +254,8 @@ def test_deep_association_with_deferred_parent_expiration
     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)
+  ensure
+    Thread.current[:deferred_parent_expiration] = nil
   end
 
   private