From 0506e8d55b508164aa6b377e9dab2c64309eae34 Mon Sep 17 00:00:00 2001
From: Donal McBreen <donal@basecamp.com>
Date: Wed, 13 Sep 2023 10:21:18 +0100
Subject: [PATCH] Setup the cache in an initializer

The cache cannot be setup until ActiveRecord has been initialized. We'll
configure it in an after_initialize hook to avoid the hit of building
the consistent hash on the first cache lookup.
---
 .../cluster/connection_handling.rb            | 50 ++++++++++---------
 lib/solid_cache/engine.rb                     |  4 ++
 lib/solid_cache/store.rb                      |  4 ++
 test/unit/cluster_test.rb                     |  1 +
 4 files changed, 35 insertions(+), 24 deletions(-)

diff --git a/lib/solid_cache/cluster/connection_handling.rb b/lib/solid_cache/cluster/connection_handling.rb
index 762026f..895d1c2 100644
--- a/lib/solid_cache/cluster/connection_handling.rb
+++ b/lib/solid_cache/cluster/connection_handling.rb
@@ -15,17 +15,40 @@ def initialize(options = {})
         end
 
         # Done lazily as the cache maybe created before ActionRecord initialization
-        @shards_initialized = false
+        @setup = false
+      end
+
+      def setup?
+        @setup
+      end
+
+      def setup!
+        return if setup?
+
+        case @shard_options
+        when Array, NilClass
+          @shards = @shard_options || SolidCache.all_shard_keys || []
+          @nodes = @shards.to_h { |shard| [ shard, shard ] }
+        when Hash
+          @shards = @shard_options.keys
+          @nodes = @shard_options.invert
+        end
+
+        if @shards.count > 1
+          @consistent_hash = MaglevHash.new(@nodes.keys)
+        end
+
+        @setup = true
       end
 
       def shards
-        initialize_shards unless shards_initialized?
+        setup!
 
         @shards
       end
 
       def nodes
-        initialize_shards unless shards_initialized?
+        setup!
 
         @nodes
       end
@@ -71,27 +94,6 @@ def reading_shard(normalized_key:)
       private
         attr_reader :consistent_hash
 
-        def shards_initialized?
-          @shards_initialized
-        end
-
-        def initialize_shards
-          case @shard_options
-          when Array, NilClass
-            @shards = @shard_options || SolidCache.all_shard_keys || []
-            @nodes = @shards.to_h { |shard| [ shard, shard ] }
-          when Hash
-            @shards = @shard_options.keys
-            @nodes = @shard_options.invert
-          end
-
-          if @shards.count > 1
-            @consistent_hash = MaglevHash.new(@nodes.keys)
-          end
-
-          @shards_initialized = true
-        end
-
         def with_shard(shard)
           if shard
             Record.connected_to(shard: shard) { yield }
diff --git a/lib/solid_cache/engine.rb b/lib/solid_cache/engine.rb
index 9865a90..7c76535 100644
--- a/lib/solid_cache/engine.rb
+++ b/lib/solid_cache/engine.rb
@@ -13,5 +13,9 @@ class Engine < ::Rails::Engine
       SolidCache.executor = config.solid_cache.executor
       SolidCache.connects_to = config.solid_cache.connects_to
     end
+
+    config.after_initialize do
+      Rails.cache.setup! if Rails.cache.is_a?(Store)
+    end
   end
 end
diff --git a/lib/solid_cache/store.rb b/lib/solid_cache/store.rb
index 3463958..ea4a464 100644
--- a/lib/solid_cache/store.rb
+++ b/lib/solid_cache/store.rb
@@ -34,6 +34,10 @@ def initialize(options = {})
       @primary_cluster = clusters.first
     end
 
+    def setup!
+      clusters.each(&:setup!)
+    end
+
     def delete_matched(matcher, options = {})
       instrument :delete_matched, matcher do
         raise ArgumentError, "Only strings are supported: #{matcher.inspect}" unless String === matcher
diff --git a/test/unit/cluster_test.rb b/test/unit/cluster_test.rb
index 303054a..6dba59f 100644
--- a/test/unit/cluster_test.rb
+++ b/test/unit/cluster_test.rb
@@ -48,6 +48,7 @@ class ClusterTest < ActiveSupport::TestCase
 
     @primary_cache.delete("foo")
     @cache.fetch("foo") { 2 }
+    sleep 0.1
 
     assert_equal 2, @cache.read("foo")
     assert_equal 2, @primary_cache.read("foo")