redis-memo 0.0.0.beta.1 → 0.0.0.beta.6

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 3e64eaaecf192519453c1accb0d10ca25ed2cf8274c10b6cb128a14e9a2e616e
-  data.tar.gz: fd98fce04f77660b8fc5960e491695504ef438394a882fb80d36b706038d10e5
+  metadata.gz: 10f721bc9b55f26c51674ce16f120bc643c306748dd9ff5aa5f54a7d9b944b51
+  data.tar.gz: 3cf3fefe9a5119d231bf973da22aa7787fcac1cf71f46ab0cadc5421cc82da52
 SHA512:
-  metadata.gz: 20ca8f3efa4ca8591d073960c5fd5b9c8fbcd920b58ae6fe378ec3511af082c2cf964c72cda93739b6bef7294d5be9ceb77eb6000cb3df567f18b2bd581cf4c8
-  data.tar.gz: 1c485024e2eb3f541139bbd8f136279f5b1010b7b81ab30929d2e87a8917c4f0d8fb35d2ae39d94fd14d2f918b6476f53d589026f4412a8b118d6b2b6d732a93
+  metadata.gz: a576fa80cef7462769055bdf657e0ee0f70fc9fb63a08f5b4ceabb931faaa8c6c8d310f01424c2124c7842b00d3c069d7c31499fc4dff166544e85b52989d0bd
+  data.tar.gz: 45cc8638df93f0c22fd050ad39b25fd1d8671ac309fbee8e41979d51e0a6136e7c98d9fc9a31acd9e1e1828c592f5b2068349f6a0e1bb294664fcacea8911f48

data/lib/redis-memo.rb

@@ -1 +1,5 @@
+# frozen_string_literal: true
+
+# Automatically require the main file of the redis-memo gem when adding it to
+# the Gemfile
 require_relative 'redis_memo'

data/lib/redis_memo.rb

@@ -1,18 +1,53 @@
 # frozen_string_literal: true
+
 require 'active_support/all'
 require 'digest'
 require 'json'
+require 'securerandom'
 
 module RedisMemo
   require 'redis_memo/memoize_method'
   require 'redis_memo/memoize_query'
 
+  # A process-level +RedisMemo::Options+ instance that stores the global
+  # options. This object can be modified by +RedisMemo.configure+.
+  #
+  # +memoize_method+ allows users to provide method-level configuration.
+  # When no callsite-level configuration specified we will use the values in
+  # +DefaultOptions+ as the default value.
   DefaultOptions = RedisMemo::Options.new
 
+  # @todo Move thread keys to +RedisMemo::ThreadKey+
+  THREAD_KEY_WITHOUT_MEMO = :__redis_memo_without_memo__
+
+  # Configure global-level default options. Those options will be used unless
+  # some options specified at +memoize_method+ callsite level. See
+  # +RedisMemo::Options+ for all the possible options.
+  #
+  # @yieldparam [RedisMemo::Options] default_options
+  # +RedisMemo::DefaultOptions+
+  # @return [void]
   def self.configure(&blk)
     blk.call(DefaultOptions)
   end
 
+  # Batch Redis calls triggered by +memoize_method+ to minimize the round trips
+  # to Redis.
+  # - Batches cannot be nested
+  # - When a batch is still open (while still in the +RedisMemo.batch+ block)
+  # the return value of any memoized method is a +RedisMemo::Future+ instead of
+  # the actual method value
+  # - The actual method values are returned as a list, in the same order as
+  # invoking, after exiting the block
+  #
+  # @example
+  #   results = RedisMemo.batch do
+  #     5.times { |i| memoized_calculation(i) }
+  #     nil # Not the return value of the block
+  #   end
+  #   results.size == 5 # true
+  #
+  # See +RedisMemo::Batch+ for more information.
   def self.batch(&blk)
     RedisMemo::Batch.open
     blk.call
@@ -21,10 +56,17 @@ module RedisMemo
     RedisMemo::Batch.close
   end
 
+  # @todo Move this method out of the top namespace
   def self.checksum(serialized)
     Digest::SHA1.base64digest(serialized)
   end
 
+  # @todo Move this method out of the top namespace
+  def self.uuid
+    SecureRandom.uuid
+  end
+
+  # @todo Move this method out of the top namespace
   def self.deep_sort_hash(orig_hash)
     {}.tap do |new_hash|
       orig_hash.sort.each do |k, v|
@@ -33,12 +75,17 @@ module RedisMemo
     end
   end
 
-  THREAD_KEY_WITHOUT_MEMO = :__redis_memo_without_memo__
-
+  # Whether the current execution context has been configured to skip
+  # memoization and use the uncached code path.
+  #
+  # @return [Boolean]
   def self.without_memo?
     Thread.current[THREAD_KEY_WITHOUT_MEMO] == true
   end
 
+  # Configure the wrapped code in the block to skip memoization.
+  #
+  # @yield [] no_args The block of code to skip memoization.
   def self.without_memo
     prev_value = Thread.current[THREAD_KEY_WITHOUT_MEMO]
     Thread.current[THREAD_KEY_WITHOUT_MEMO] = true
@@ -47,6 +94,8 @@ module RedisMemo
     Thread.current[THREAD_KEY_WITHOUT_MEMO] = prev_value
   end
 
+  # @todo Move errors to a separate file errors.rb
   class ArgumentError < ::ArgumentError; end
   class RuntimeError < ::RuntimeError; end
+  class WithoutMemoization < Exception; end
 end

data/lib/redis_memo/after_commit.rb

@@ -66,7 +66,12 @@ class RedisMemo::AfterCommit
       @@pending_memo_versions.each do |key, version|
         invalidation_queue =
           RedisMemo::Memoizable::Invalidation.class_variable_get(:@@invalidation_queue)
-        invalidation_queue << [key, version, @@previous_memo_versions[key]]
+
+        invalidation_queue << RedisMemo::Memoizable::Invalidation::Task.new(
+          key,
+          version,
+          @@previous_memo_versions[key],
+        )
       end
 
       RedisMemo::Memoizable::Invalidation.drain_invalidation_queue

data/lib/redis_memo/cache.rb

@@ -1,12 +1,14 @@
  # frozen_string_literal: true
 require_relative 'options'
 require_relative 'redis'
+require_relative 'connection_pool'
 
 class RedisMemo::Cache < ActiveSupport::Cache::RedisCacheStore
   class Rescuable < Exception; end
 
-  THREAD_KEY_LOCAL_CACHE = :__redis_memo_cache_local_cache__
-  THREAD_KEY_RAISE_ERROR = :__redis_memo_cache_raise_error__
+  THREAD_KEY_LOCAL_CACHE            = :__redis_memo_cache_local_cache__
+  THREAD_KEY_LOCAL_DEPENDENCY_CACHE = :__redis_memo_local_cache_dependency_cache__
+  THREAD_KEY_RAISE_ERROR            = :__redis_memo_cache_raise_error__
 
   @@redis = nil
   @@redis_store = nil
@@ -23,7 +25,12 @@ class RedisMemo::Cache < ActiveSupport::Cache::RedisCacheStore
   end
 
   def self.redis
-    @@redis ||= RedisMemo::DefaultOptions.redis
+    @@redis ||=
+      if RedisMemo::DefaultOptions.connection_pool
+        RedisMemo::ConnectionPool.new(**RedisMemo::DefaultOptions.connection_pool)
+      else
+        RedisMemo::DefaultOptions.redis
+      end
   end
 
   def self.redis_store
@@ -44,16 +51,22 @@ class RedisMemo::Cache < ActiveSupport::Cache::RedisCacheStore
     Thread.current[THREAD_KEY_LOCAL_CACHE]
   end
 
+  def self.local_dependency_cache
+    Thread.current[THREAD_KEY_LOCAL_DEPENDENCY_CACHE]
+  end
+
   class << self
     def with_local_cache(&blk)
       Thread.current[THREAD_KEY_LOCAL_CACHE] = {}
+      Thread.current[THREAD_KEY_LOCAL_DEPENDENCY_CACHE] = {}
       blk.call
     ensure
       Thread.current[THREAD_KEY_LOCAL_CACHE] = nil
+      Thread.current[THREAD_KEY_LOCAL_DEPENDENCY_CACHE] = nil
     end
 
     # RedisCacheStore doesn't read from the local cache before reading from redis
-    def read_multi(*keys, raise_error: false)
+    def read_multi(*keys, raw: false, raise_error: false)
       return {} if keys.empty?
 
       Thread.current[THREAD_KEY_RAISE_ERROR] = raise_error
@@ -64,7 +77,7 @@ class RedisMemo::Cache < ActiveSupport::Cache::RedisCacheStore
       keys_to_fetch -= local_entries.keys unless local_entries.empty?
       return local_entries if keys_to_fetch.empty?
 
-      remote_entries = redis_store.read_multi(*keys_to_fetch)
+      remote_entries = redis_store.read_multi(*keys_to_fetch, raw: raw)
       local_cache&.merge!(remote_entries)
 
       if local_entries.empty?

data/lib/redis_memo/connection_pool.rb

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+require 'connection_pool'
+require_relative 'redis'
+
+class RedisMemo::ConnectionPool
+  def initialize(**options)
+    @connection_pool = ::ConnectionPool.new(**options) do
+      # Construct a new client every time the block gets called
+      RedisMemo::Redis.new(RedisMemo::DefaultOptions.redis_config)
+    end
+  end
+
+  # Avoid method_missing when possible for better performance
+  %i(get mget mapped_mget set eval).each do |method_name|
+    define_method method_name do |*args, &blk|
+      @connection_pool.with do |redis|
+        redis.send(method_name, *args, &blk)
+      end
+    end
+  end
+
+  def method_missing(method_name, *args, &blk)
+    @connection_pool.with do |redis|
+      redis.send(method_name, *args, &blk)
+    end
+  end
+end

data/lib/redis_memo/future.rb

@@ -9,14 +9,14 @@ class RedisMemo::Future
     ref,
     method_id,
     method_args,
-    depends_on,
+    dependent_memos,
     cache_options,
     method_name_without_memo
   )
     @ref = ref
     @method_id = method_id
     @method_args = method_args
-    @depends_on = depends_on
+    @dependent_memos = dependent_memos
     @cache_options = cache_options
     @method_name_without_memo = method_name_without_memo
     @method_cache_key = nil
@@ -28,7 +28,7 @@ class RedisMemo::Future
   end
 
   def context
-    [@ref, @method_id, @method_args, @depends_on]
+    [@method_id, @method_args, @dependent_memos]
   end
 
   def method_cache_key

data/lib/redis_memo/memoizable.rb

@@ -1,5 +1,4 @@
 # frozen_string_literal: true
-require 'securerandom'
 
 class RedisMemo::Memoizable
   require_relative 'memoizable/dependency'
@@ -52,7 +51,7 @@ class RedisMemo::Memoizable
       cache_key = instance.cache_key
       RedisMemo::Memoizable::Invalidation.bump_version_later(
         cache_key,
-        SecureRandom.uuid,
+        RedisMemo.uuid,
       )
     end
 
@@ -83,7 +82,11 @@ class RedisMemo::Memoizable
       if keys_to_fetch.empty?
         {}
       else
-        RedisMemo::Cache.read_multi(*keys_to_fetch, raise_error: true)
+        RedisMemo::Cache.read_multi(
+          *keys_to_fetch,
+          raw: true,
+          raise_error: true,
+        )
       end
     memo_versions.merge!(cached_versions) unless cached_versions.empty?
 
@@ -96,7 +99,7 @@ class RedisMemo::Memoizable
         # cached result.
         need_to_bump_versions = true
 
-        new_version = SecureRandom.uuid
+        new_version = RedisMemo.uuid
         RedisMemo::Memoizable::Invalidation.bump_version_later(
           key,
           new_version,

data/lib/redis_memo/memoizable/dependency.rb

@@ -12,25 +12,55 @@ class RedisMemo::Memoizable::Dependency
     @nodes.values
   end
 
-  def depends_on(memo_or_model, **conditions)
-    if !memo_or_model.is_a?(RedisMemo::Memoizable)
+  def depends_on(dependency, **conditions)
+    case dependency
+    when self.class
+      nodes.merge!(dependency.nodes)
+    when RedisMemo::Memoizable
+      memo = dependency
+      return if nodes.include?(memo.cache_key)
+      nodes[memo.cache_key] = memo
+
+      if memo.depends_on
+        # Extract dependencies from the current memoizable and recurse
+        instance_exec(&memo.depends_on)
+      end
+    when ActiveRecord::Relation
+      extracted = extract_dependencies_for_relation(dependency)
+      nodes.merge!(extracted.nodes)
+    when UsingActiveRecord
       [
-        memo_or_model.redis_memo_class_memoizable,
-        RedisMemo::MemoizeQuery.create_memo(memo_or_model, **conditions),
+        dependency.redis_memo_class_memoizable,
+        RedisMemo::MemoizeQuery.create_memo(dependency, **conditions),
       ].each do |memo|
         nodes[memo.cache_key] = memo
       end
-
-      return
+    else
+      raise(
+        RedisMemo::ArgumentError,
+        "Invalid dependency #{dependency}"
+      )
     end
+  end
 
-    memo = memo_or_model
-    return if nodes.include?(memo.cache_key)
-    nodes[memo.cache_key] = memo
+  def extract_dependencies_for_relation(relation)
+    # Extract the dependent memos of an Arel without calling exec_query to actually execute the query
+    RedisMemo::MemoizeQuery::CachedSelect.with_new_query_context do
+      connection = ActiveRecord::Base.connection
+      query, binds, _ = connection.send(:to_sql_and_binds, relation.arel)
+      RedisMemo::MemoizeQuery::CachedSelect.current_query = relation.arel
+      is_query_cached = RedisMemo::MemoizeQuery::CachedSelect.extract_bind_params(query)
+        raise(
+          RedisMemo::WithoutMemoization,
+          "Arel query is not cached using RedisMemo."
+        ) unless is_query_cached
+        extracted_dependency = connection.dependency_of(:exec_query, query, nil, binds)
+    end
+  end
 
-    if memo.depends_on
-      # Extract dependencies from the current memoizable and recurse
-      instance_exec(&memo.depends_on)
+  class UsingActiveRecord
+    def self.===(dependency)
+      RedisMemo::MemoizeQuery.using_active_record?(dependency)
     end
   end
 end

data/lib/redis_memo/memoizable/invalidation.rb

@@ -3,12 +3,29 @@ require_relative '../after_commit'
 require_relative '../cache'
 
 module RedisMemo::Memoizable::Invalidation
-  # This is a thread safe data structure
-  #
-  # Handle transient network errors during cache invalidation
-  # Each item in the queue is a tuple:
-  #
-  #   [key, version (a UUID), previous_version (a UUID)]
+  class Task
+    attr_reader :key
+    attr_reader :version
+    attr_reader :previous_version
+
+    def initialize(key, version, previous_version)
+      @key = key
+      @version = version
+      @previous_version = previous_version
+      @created_at = current_timestamp
+    end
+
+    def current_timestamp
+      Process.clock_gettime(Process::CLOCK_REALTIME, :nanosecond)
+    end
+
+    def duration
+      current_timestamp - @created_at
+    end
+  end
+
+  # This is a thread safe data structure to handle transient network errors
+  # during cache invalidation
   #
   # When an invalidation call arrives at Redis, we only bump to the specified
   # version (so the cached results using that version will become visible) if
@@ -28,8 +45,6 @@ module RedisMemo::Memoizable::Invalidation
   @@invalidation_queue = Queue.new
 
   def self.bump_version_later(key, version, previous_version: nil)
-    RedisMemo::DefaultOptions.logger&.info("[received] Bump memo key #{key}")
-
     if RedisMemo::AfterCommit.in_transaction?
       previous_version ||= RedisMemo::AfterCommit.pending_memo_versions[key]
     end
@@ -41,7 +56,10 @@ module RedisMemo::Memoizable::Invalidation
       # Fill an expected previous version so the later calculation results
       # based on this version can still be rolled out if this version
       # does not change
-      previous_version ||= RedisMemo::Cache.read_multi(key)[key]
+      previous_version ||= RedisMemo::Cache.read_multi(
+        key,
+        raw: true,
+      )[key]
     end
 
     local_cache&.send(:[]=, key, version)
@@ -52,7 +70,7 @@ module RedisMemo::Memoizable::Invalidation
         previous_version: previous_version,
       )
     else
-      @@invalidation_queue << [key, version, previous_version]
+      @@invalidation_queue << Task.new(key, version, previous_version)
     end
   end
 
@@ -89,31 +107,32 @@ module RedisMemo::Memoizable::Invalidation
     return redis.call('set', key, new_version, unpack(px))
   LUA
 
-  def self.bump_version(cache_key, version, previous_version:)
-    RedisMemo::Tracer.trace('redis_memo.memoizable.bump_version', nil) do
+  def self.bump_version(task)
+    RedisMemo::Tracer.trace('redis_memo.memoizable.bump_version', task.key) do
       ttl = RedisMemo::DefaultOptions.expires_in
       ttl = (ttl * 1000.0).to_i if ttl
       RedisMemo::Cache.redis.eval(
         LUA_BUMP_VERSION,
-        keys: [cache_key],
-        argv: [previous_version, version, SecureRandom.uuid, ttl],
+        keys: [task.key],
+        argv: [task.previous_version, task.version, RedisMemo.uuid, ttl],
       )
+      RedisMemo::Tracer.set_tag(enqueue_to_finish: task.duration)
     end
-    RedisMemo::DefaultOptions.logger&.info("[performed] Bump memo key #{cache_key}")
   end
 
   def self.drain_invalidation_queue_now
     retry_queue = []
     until @@invalidation_queue.empty?
-      tuple = @@invalidation_queue.pop
+      task = @@invalidation_queue.pop
       begin
-        bump_version(tuple[0], tuple[1], previous_version: tuple[2])
-      rescue SignalException, Redis::BaseConnectionError
-        retry_queue << tuple
+        bump_version(task)
+      rescue SignalException, Redis::BaseConnectionError,
+             ::ConnectionPool::TimeoutError
+        retry_queue << task
       end
     end
   ensure
-    retry_queue.each { |t| @@invalidation_queue << t }
+    retry_queue.each { |task| @@invalidation_queue << task }
   end
 
   at_exit do

data/lib/redis_memo/memoize_method.rb

@@ -15,6 +15,12 @@ module RedisMemo::MemoizeMethod
     define_method method_name_with_memo do |*args|
       return send(method_name_without_memo, *args) if RedisMemo.without_memo?
 
+      dependent_memos = nil
+      if depends_on
+        dependency = RedisMemo::MemoizeMethod.get_or_extract_dependencies(self, *args, &depends_on)
+        dependent_memos = dependency.memos
+      end
+
       future = RedisMemo::Future.new(
         self,
         case method_id
@@ -26,7 +32,7 @@ module RedisMemo::MemoizeMethod
           method_id.call(self, *args)
         end,
         args,
-        depends_on,
+        dependent_memos,
         options,
         method_name_without_memo,
       )
@@ -37,9 +43,25 @@ module RedisMemo::MemoizeMethod
       end
 
       future.execute
+    rescue RedisMemo::WithoutMemoization
+      send(method_name_without_memo, *args)
     end
 
     alias_method method_name, method_name_with_memo
+
+    @__redis_memo_method_dependencies ||= Hash.new
+    @__redis_memo_method_dependencies[method_name] = depends_on
+
+    define_method :dependency_of do |method_name, *method_args|
+      method_depends_on = self.class.instance_variable_get(:@__redis_memo_method_dependencies)[method_name]
+      unless method_depends_on
+        raise(
+          RedisMemo::ArgumentError,
+          "#{method_name} is not a memoized method"
+        )
+      end
+      RedisMemo::MemoizeMethod.get_or_extract_dependencies(self, *method_args, &method_depends_on)
+    end
   end
 
   def self.method_id(ref, method_name)
@@ -49,23 +71,34 @@ module RedisMemo::MemoizeMethod
     "#{class_name}#{is_class_method ? '::' : '#'}#{method_name}"
   end
 
-  def self.method_cache_keys(future_contexts)
-    memos = Array.new(future_contexts.size)
-    future_contexts.each_with_index do |(ref, _, method_args, depends_on), i|
-      if depends_on
-        dependency = RedisMemo::Memoizable::Dependency.new
+  def self.extract_dependencies(ref, *method_args, &depends_on)
+    dependency = RedisMemo::Memoizable::Dependency.new
 
-        # Resolve the dependency recursively
-        dependency.instance_exec(ref, *method_args, &depends_on)
+    # Resolve the dependency recursively
+    dependency.instance_exec(ref, *method_args, &depends_on)
+    dependency
+  end
 
-        memos[i] = dependency.memos
-      end
+  def self.get_or_extract_dependencies(ref, *method_args, &depends_on)
+    if RedisMemo::Cache.local_dependency_cache
+      RedisMemo::Cache.local_dependency_cache[ref] ||= {}
+      RedisMemo::Cache.local_dependency_cache[ref][depends_on] ||= {}
+      RedisMemo::Cache.local_dependency_cache[ref][depends_on][method_args] ||= extract_dependencies(ref, *method_args, &depends_on)
+    else
+      extract_dependencies(ref, *method_args, &depends_on)
+    end
+  end
+
+  def self.method_cache_keys(future_contexts)
+    memos = Array.new(future_contexts.size)
+    future_contexts.each_with_index do |(_, _, dependent_memos), i|
+      memos[i] = dependent_memos
     end
 
     j = 0
     memo_checksums = RedisMemo::Memoizable.checksums(memos.compact)
     method_cache_key_versions = Array.new(future_contexts.size)
-    future_contexts.each_with_index do |(_, method_id, method_args, _), i|
+    future_contexts.each_with_index do |(method_id, method_args, _), i|
       if memos[i]
         method_cache_key_versions[i] = [method_id, memo_checksums[j]]
         j += 1

data/lib/redis_memo/memoize_query.rb

@@ -13,6 +13,8 @@ if defined?(ActiveRecord)
     # after each record save
     def memoize_table_column(*raw_columns, editable: true)
       RedisMemo::MemoizeQuery.using_active_record!(self)
+      return if ENV["REDIS_MEMO_DISABLE_#{self.table_name.upcase}"] == 'true'
+
       columns = raw_columns.map(&:to_sym).sort
 
       RedisMemo::MemoizeQuery.memoized_columns(self, editable_only: true) << columns if editable
@@ -44,11 +46,15 @@ if defined?(ActiveRecord)
     end
 
     def self.using_active_record!(model_class)
-      unless model_class.respond_to?(:<) && model_class < ActiveRecord::Base
+      unless using_active_record?(model_class)
         raise RedisMemo::ArgumentError, "'#{model_class.name}' does not use ActiveRecord"
       end
     end
 
+    def self.using_active_record?(model_class)
+      model_class.respond_to?(:<) && model_class < ActiveRecord::Base
+    end
+
     @@memoized_columns = Hash.new { |h, k| h[k] = [Set.new, Set.new] }
 
     def self.memoized_columns(model_or_table, editable_only: false)
@@ -68,26 +74,23 @@ if defined?(ActiveRecord)
         )
       end
 
-      extra_props.each do |key, values|
+      extra_props.each do |key, value|
         # The data type is ensured by the database, thus we don't need to cast
         # types here for better performance
         column_name = key.to_s
-        values = [values] unless values.is_a?(Enumerable)
         extra_props[key] =
           if model_class.defined_enums.include?(column_name)
             enum_mapping = model_class.defined_enums[column_name]
-            values.map do |value|
-              # Assume a value is a converted enum if it does not exist in the
-              # enum mapping
-              (enum_mapping[value.to_s] || value).to_s
-            end
+            # Assume a value is a converted enum if it does not exist in the
+            # enum mapping
+            (enum_mapping[value.to_s] || value).to_s
           else
-            values.map(&:to_s)
+            value.to_s
           end
       end
 
       RedisMemo::Memoizable.new(
-        __redis_memo_memoize_query_model_class_name__: model_class.name,
+        __redis_memo_memoize_query_table_name__: model_class.table_name,
         **extra_props,
       )
     end
@@ -101,7 +104,13 @@ if defined?(ActiveRecord)
       end
     end
 
-    def self.invalidate(record)
+    def self.invalidate(*records)
+      RedisMemo::Memoizable.invalidate(
+        records.map { |record| to_memos(record) }.flatten,
+      )
+    end
+
+    def self.to_memos(record)
       # Invalidate memos with current values
       memos_to_invalidate = memoized_columns(record.class).map do |columns|
         props = {}
@@ -136,7 +145,7 @@ if defined?(ActiveRecord)
         end
       end
 
-      RedisMemo::Memoizable.invalidate(memos_to_invalidate)
+      memos_to_invalidate
     end
   end
 end

data/lib/redis_memo/memoize_query/cached_select.rb

@@ -123,7 +123,14 @@ class RedisMemo::MemoizeQuery::CachedSelect
 
         depends_on RedisMemo::Memoizable.new(
           __redis_memo_memoize_query_memoize_query_sql__: sql,
-          __redis_memo_memoize_query_memoize_query_binds__: binds.map(&:value_for_database),
+          __redis_memo_memoize_query_memoize_query_binds__: binds.map do |bind|
+            if bind.respond_to?(:value_for_database)
+              bind.value_for_database
+            else
+              # In activerecord >= 6, a bind could be an actual database value
+              bind
+            end
+          end
         )
       end
     end
@@ -187,6 +194,19 @@ class RedisMemo::MemoizeQuery::CachedSelect
     Thread.current[THREAD_KEY_AREL_BIND_PARAMS] = nil
   end
 
+  def self.with_new_query_context
+    prev_arel = Thread.current[THREAD_KEY_AREL]
+    prev_substitutes = Thread.current[THREAD_KEY_SUBSTITUTES]
+    prev_bind_params = Thread.current[THREAD_KEY_AREL_BIND_PARAMS]
+    RedisMemo::MemoizeQuery::CachedSelect.reset_current_query
+
+    yield
+  ensure
+    Thread.current[THREAD_KEY_AREL] = prev_arel
+    Thread.current[THREAD_KEY_SUBSTITUTES] = prev_substitutes
+    Thread.current[THREAD_KEY_AREL_BIND_PARAMS] = prev_bind_params
+  end
+
   private
 
   # A pre-order Depth First Search
@@ -197,7 +217,7 @@ class RedisMemo::MemoizeQuery::CachedSelect
     bind_params = BindParams.new
 
     case node
-    when Arel::Nodes::Equality, Arel::Nodes::In
+    when NodeHasFilterCondition
       attr_node = node.left
       return unless attr_node.is_a?(Arel::Attributes::Attribute)
 
@@ -234,7 +254,13 @@ class RedisMemo::MemoizeQuery::CachedSelect
           }
         when Arel::Nodes::Casted
           bind_params.params[binding_relation] << {
-            right.attribute.name.to_sym => right.val,
+            right.attribute.name.to_sym =>
+              if right.respond_to?(:val)
+                right.val
+              else
+                # activerecord >= 6
+                right.value
+              end,
           }
         else
           bind_params = bind_params.union(extract_bind_params_recurse(right))
@@ -330,6 +356,23 @@ class RedisMemo::MemoizeQuery::CachedSelect
     enabled_models[table_node.try(:name)]
   end
 
+  class NodeHasFilterCondition
+    def self.===(node)
+      case node
+      when Arel::Nodes::Equality, Arel::Nodes::In
+        true
+      else
+        # In activerecord >= 6, a new arel node HomogeneousIn is introduced
+        if defined?(Arel::Nodes::HomogeneousIn) &&
+           node.is_a?(Arel::Nodes::HomogeneousIn)
+          true
+        else
+          false
+        end
+      end
+    end
+  end
+
   # Thread locals to exchange information between RedisMemo and ActiveRecord
   THREAD_KEY_AREL = :__redis_memo_memoize_query_cached_select_arel__
   THREAD_KEY_SUBSTITUTES = :__redis_memo_memoize_query_cached_select_substitues__

data/lib/redis_memo/memoize_query/invalidation.rb

@@ -33,7 +33,6 @@ class RedisMemo::MemoizeQuery::Invalidation
     # Methods that won't trigger model callbacks
     # https://guides.rubyonrails.org/active_record_callbacks.html#skipping-callbacks
     %i(
-      import
       decrement_counter
       delete_all delete_by
       increment_counter
@@ -43,7 +42,7 @@ class RedisMemo::MemoizeQuery::Invalidation
       upsert upsert_all
     ).each do |method_name|
       # Example: Model.update_all
-      rewrite_bulk_update_method(
+      rewrite_default_method(
         model_class,
         model_class,
         method_name,
@@ -51,7 +50,7 @@ class RedisMemo::MemoizeQuery::Invalidation
       )
 
       # Example: Model.where(...).update_all
-      rewrite_bulk_update_method(
+      rewrite_default_method(
         model_class,
         model_class.const_get(:ActiveRecord_Relation),
         method_name,
@@ -59,20 +58,28 @@ class RedisMemo::MemoizeQuery::Invalidation
       )
     end
 
+    %i(
+      import import!
+    ).each do |method_name|
+      rewrite_import_method(
+        model_class,
+        method_name,
+      )
+    end
+
     model_class.class_variable_set(var_name, true)
   end
 
   private
 
   #
-  # There’s no good way to perform fine-grind cache invalidation when operations
-  # are bulk update operations such as import, update_all, and destroy_all:
-  # Performing fine-grind cache invalidation would require the applications to
-  # fetch additional data from the database, which might lead to performance
-  # degradation. Thus we simply invalidate all existing cached records after each
-  # bulk_updates.
+  # There’s no good way to perform fine-grind cache invalidation when
+  # operations are bulk update operations such as update_all, and delete_all
+  # witout fetching additional data from the database, which might lead to
+  # performance degradation. Thus, by default, we simply invalidate all
+  # existing cached records after each bulk_updates.
   #
-  def self.rewrite_bulk_update_method(model_class, klass, method_name, class_method:)
+  def self.rewrite_default_method(model_class, klass, method_name, class_method:)
     methods = class_method ? :methods : :instance_methods
     return unless klass.send(methods).include?(method_name)
 
@@ -87,4 +94,120 @@ class RedisMemo::MemoizeQuery::Invalidation
       end
     end
   end
+
+  def self.rewrite_import_method(model_class, method_name)
+    # This optimization to avoid over-invalidation only works on postgres
+    unless ActiveRecord::Base.connection.is_a?(ActiveRecord::ConnectionAdapters::PostgreSQLAdapter)
+      rewrite_default_method(model_class, model_class, method_name, class_method: true)
+      return
+    end
+
+    model_class.singleton_class.class_eval do
+      alias_method :"#{method_name}_without_redis_memo_invalidation", method_name
+
+      # For the args format, see
+      # https://github.com/zdennis/activerecord-import/blob/master/lib/activerecord-import/import.rb#L128
+      define_method method_name do |*args, &blk|
+        options = args.last.is_a?(Hash) ? args.last : {}
+        records = args[args.last.is_a?(Hash) ? -2 : -1]
+        columns_to_update = options[:on_duplicate_key_update]
+        if columns_to_update.is_a?(Hash)
+          columns_to_update = columns_to_update[:columns]
+        end
+
+        if records.last.is_a?(Hash)
+          records.map! { |hash| model_class.new(hash) }
+        end
+
+        # Invalidate the records before and after the import to resolve
+        # - default values filled by the database
+        # - updates on conflict conditions
+        records_to_invalidate =
+          if columns_to_update
+            RedisMemo::MemoizeQuery::Invalidation.send(
+              :select_by_columns,
+              model_class,
+              records,
+              columns_to_update,
+            )
+          else
+            []
+          end
+
+        result = send(:"#{method_name}_without_redis_memo_invalidation", *args, &blk)
+
+        # Offload the records to invalidate while selecting the next set of
+        # records to invalidate
+        case records_to_invalidate
+        when Array
+          RedisMemo::MemoizeQuery.invalidate(*records_to_invalidate) unless records_to_invalidate.empty?
+
+          RedisMemo::MemoizeQuery.invalidate(*RedisMemo::MemoizeQuery::Invalidation.send(
+            :select_by_id,
+            model_class,
+            # Not all databases support "RETURNING", which is useful when
+            # invaldating records after bulk creation
+            result.ids,
+          ))
+        else
+          RedisMemo::MemoizeQuery.invalidate_all(model_class)
+        end
+
+        result
+      end
+    end
+  end
+
+  def self.select_by_columns(model_class, records, columns_to_update)
+    return [] if records.empty?
+
+    or_chain = nil
+    columns_to_select = columns_to_update & RedisMemo::MemoizeQuery
+      .memoized_columns(model_class)
+      .to_a.flatten.uniq
+
+    # Nothing to invalidate here
+    return [] if columns_to_select.empty?
+
+    RedisMemo::Tracer.trace(
+      'redis_memo.memoize_query.invalidation',
+      "#{__method__}##{model_class.name}",
+    ) do
+      records.each do |record|
+        conditions = {}
+        columns_to_select.each do |column|
+          conditions[column] = record.send(column)
+        end
+        if or_chain
+          or_chain = or_chain.or(model_class.where(conditions))
+        else
+          or_chain = model_class.where(conditions)
+        end
+      end
+
+      record_count = RedisMemo.without_memo { or_chain.count }
+      if record_count > bulk_operations_invalidation_limit
+        nil
+      else
+        RedisMemo.without_memo { or_chain.to_a }
+      end
+    end
+  end
+
+  def self.select_by_id(model_class, ids)
+    RedisMemo::Tracer.trace(
+      'redis_memo.memoize_query.invalidation',
+      "#{__method__}##{model_class.name}",
+    ) do
+      RedisMemo.without_memo do
+        model_class.where(model_class.primary_key => ids).to_a
+      end
+    end
+  end
+
+  def self.bulk_operations_invalidation_limit
+    ENV['REDIS_MEMO_BULK_OPERATIONS_INVALIDATION_LIMIT']&.to_i ||
+      RedisMemo::DefaultOptions.bulk_operations_invalidation_limit ||
+      10000
+  end
 end

data/lib/redis_memo/options.rb

@@ -13,7 +13,7 @@ class RedisMemo::Options
   )
     @compress = compress.nil? ? true : compress
     @compress_threshold = compress_threshold || 1.kilobyte
-    @redis = redis
+    @redis_config = redis
     @redis_client = nil
     @redis_error_handler = redis_error_handler
     @tracer = tracer
@@ -22,20 +22,18 @@ class RedisMemo::Options
     @expires_in = expires_in
   end
 
-  def redis(&blk)
-    if blk.nil?
-      return @redis_client if @redis_client.is_a?(RedisMemo::Redis)
+  def redis
+    @redis_client ||= RedisMemo::Redis.new(redis_config)
+  end
 
-      if @redis.respond_to?(:call)
-        @redis_client = RedisMemo::Redis.new(@redis.call)
-      elsif @redis
-        @redis_client = RedisMemo::Redis.new(@redis)
-      else
-        @redis_client = RedisMemo::Redis.new
-      end
-    else
-      @redis = blk
-    end
+  def redis_config
+    @redis_config || {}
+  end
+
+  def redis=(config)
+    @redis_config = config
+    @redis_client = nil
+    redis
   end
 
   def tracer(&blk)
@@ -74,15 +72,16 @@ class RedisMemo::Options
   end
 
   attr_accessor :async
+  attr_accessor :bulk_operations_invalidation_limit
+  attr_accessor :cache_out_of_date_handler
+  attr_accessor :cache_validation_sampler
   attr_accessor :compress
   attr_accessor :compress_threshold
-  attr_accessor :redis_error_handler
+  attr_accessor :connection_pool
   attr_accessor :expires_in
-  attr_accessor :cache_validation_sampler
-  attr_accessor :cache_out_of_date_handler
+  attr_accessor :redis_error_handler
 
   attr_writer :global_cache_key_version
-  attr_writer :redis
   attr_writer :tracer
   attr_writer :logger
 end

data/lib/redis_memo/redis.rb

@@ -31,7 +31,8 @@ class RedisMemo::Redis < Redis::Distributed
   end
 
   class WithReplicas < ::Redis
-    def initialize(options)
+    def initialize(orig_options)
+      options = orig_options.dup
       primary_option = options.shift
       @replicas = options.map do |option|
         option[:logger] ||= RedisMemo::DefaultOptions.logger

data/lib/redis_memo/tracer.rb

@@ -11,13 +11,15 @@ class RedisMemo::Tracer
     end
   end
 
-  def self.set_tag(cache_hit:)
+  def self.set_tag(**tags)
     tracer = RedisMemo::DefaultOptions.tracer
     return if tracer.nil? || !tracer.respond_to?(:active_span)
 
     active_span = tracer.active_span
     return if !active_span.respond_to?(:set_tag)
 
-    active_span.set_tag('cache_hit', cache_hit)
+    tags.each do |name, value|
+      active_span.set_tag(name, value)
+    end
   end
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: redis-memo
 version: !ruby/object:Gem::Version
-  version: 0.0.0.beta.1
+  version: 0.0.0.beta.6
 platform: ruby
 authors:
 - Chan Zuckerberg Initiative
@@ -14,44 +14,72 @@ dependencies:
   name: activesupport
   requirement: !ruby/object:Gem::Requirement
     requirements:
-    - - "~>"
+    - - ">="
       - !ruby/object:Gem::Version
         version: '5.2'
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
-    - - "~>"
+    - - ">="
       - !ruby/object:Gem::Version
         version: '5.2'
 - !ruby/object:Gem::Dependency
   name: redis
   requirement: !ruby/object:Gem::Requirement
     requirements:
-    - - "~>"
+    - - ">="
       - !ruby/object:Gem::Version
-        version: '4'
+        version: 4.0.1
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
-    - - "~>"
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: 4.0.1
+- !ruby/object:Gem::Dependency
+  name: connection_pool
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: 2.2.3
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
       - !ruby/object:Gem::Version
-        version: '4'
+        version: 2.2.3
 - !ruby/object:Gem::Dependency
   name: activerecord
   requirement: !ruby/object:Gem::Requirement
     requirements:
-    - - "~>"
+    - - ">="
       - !ruby/object:Gem::Version
         version: '5.2'
   type: :development
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
-    - - "~>"
+    - - ">="
       - !ruby/object:Gem::Version
         version: '5.2'
+- !ruby/object:Gem::Dependency
+  name: activerecord-import
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
 - !ruby/object:Gem::Dependency
   name: codecov
   requirement: !ruby/object:Gem::Requirement
@@ -147,6 +175,7 @@ files:
 - lib/redis_memo/after_commit.rb
 - lib/redis_memo/batch.rb
 - lib/redis_memo/cache.rb
+- lib/redis_memo/connection_pool.rb
 - lib/redis_memo/future.rb
 - lib/redis_memo/memoizable.rb
 - lib/redis_memo/memoizable/dependency.rb