Batch multi queries

app/models/solid_cache/entry.rb

@@ -14,16 +14,20 @@ class Entry < Record
 
     KEY_HASH_ID_RANGE = -(2**63)..(2**63 - 1)
 
+    MULTI_BATCH_SIZE = 1000
+
     class << self
       def write(key, value)
         write_multi([ { key: key, value: value } ])
       end
 
       def write_multi(payloads)
         without_query_cache do
-          upsert_all \
-            add_key_hash_and_byte_size(payloads),
-            unique_by: upsert_unique_by, on_duplicate: :update, update_only: [ :key, :value, :byte_size ]
+          payloads.each_slice(MULTI_BATCH_SIZE).each do |payload_batch|
+            upsert_all \
+              add_key_hash_and_byte_size(payload_batch),
+              unique_by: upsert_unique_by, on_duplicate: :update, update_only: [ :key, :value, :byte_size ]
+          end
         end
       end
 
@@ -33,9 +37,13 @@ def read(key)
 
       def read_multi(keys)
         without_query_cache do
-          query = Arel.sql(select_sql(keys), *key_hashes_for(keys))
+          {}.tap do |results|
+            keys.each_slice(MULTI_BATCH_SIZE).each do |keys_batch|
+              query = Arel.sql(select_sql(keys_batch), *key_hashes_for(keys_batch))
 
-          connection.select_all(query, "SolidCache::Entry Load").cast_values(attribute_types).to_h
+              results.merge!(connection.select_all(query, "SolidCache::Entry Load").cast_values(attribute_types).to_h)
+            end
+          end
         end
       end
 

test/models/solid_cache/entry_test.rb

@@ -70,9 +70,33 @@ class EntryTest < ActiveSupport::TestCase
       end
     end
 
+    test "batching multi queries" do
+      with_multi_batch_size(2) do
+        Entry.stubs(:const_get).with(:MULTI_BATCH_SIZE).returns("stubbed_value")
+
+        assert_queries_count(2) do
+          Entry.write_multi([ { key: "hello".b, value: "there" }, { key: "foo".b, value: "bar" }, { key: "baz".b, value: "zab" } ])
+        end
+
+        assert_queries_count(2) do
+          assert_equal({ "foo" => "bar", "hello" => "there", "baz" => "zab" }, Entry.read_multi([ "hello".b, "foo".b, "baz".b, "bar".b ]))
+        end
+      end
+    end
+
     private
       def write_entries(count = 20)
         Entry.write_multi(count.times.map { |i| { key: "key#{i}", value: "value#{i}" } })
       end
+
+      def with_multi_batch_size(value)
+        old_value = Entry::MULTI_BATCH_SIZE
+        Entry.send(:remove_const, :MULTI_BATCH_SIZE)
+        Entry.const_set(:MULTI_BATCH_SIZE, value)
+        yield
+      ensure
+        Entry.send(:remove_const, :MULTI_BATCH_SIZE)
+        Entry.const_set(:MULTI_BATCH_SIZE, old_value)
+      end
   end
 end