redis-memo 0.0.0.alpha → 0.0.0.beta.4

data/lib/redis_memo/memoizable.rb

@@ -0,0 +1,119 @@
+# frozen_string_literal: true
+
+class RedisMemo::Memoizable
+  require_relative 'memoizable/dependency'
+  require_relative 'memoizable/invalidation'
+
+  attr_accessor :props
+  attr_reader :depends_on
+
+  def initialize(**props, &depends_on)
+    @props = props
+    @depends_on = depends_on
+    @cache_key = nil
+  end
+
+  def extra_props(**args)
+    instance = dup
+    instance.props = props.dup.merge(**args)
+    instance
+  end
+
+  def cache_key
+    @cache_key ||= [
+      self.class.name,
+      RedisMemo.checksum(
+        RedisMemo.deep_sort_hash(@props).to_json,
+      ),
+    ].join(':')
+  end
+
+  # Calculate the checksums for all memoizable groups in one Redis round trip
+  def self.checksums(instances_groups)
+    dependents_cache_keys = []
+    cache_key_groups = instances_groups.map do |instances|
+      cache_keys = instances.map(&:cache_key)
+      dependents_cache_keys += cache_keys
+      cache_keys
+    end
+
+    dependents_cache_keys.uniq!
+    dependents_versions = find_or_create_versions(dependents_cache_keys)
+    version_hash = dependents_cache_keys.zip(dependents_versions).to_h
+
+    cache_key_groups.map do |cache_keys|
+      RedisMemo.checksum(version_hash.slice(*cache_keys).to_json)
+    end
+  end
+
+  def self.invalidate(instances)
+    instances.each do |instance|
+      cache_key = instance.cache_key
+      RedisMemo::Memoizable::Invalidation.bump_version_later(
+        cache_key,
+        RedisMemo.uuid,
+      )
+    end
+
+    RedisMemo::Memoizable::Invalidation.drain_invalidation_queue
+  end
+
+  private
+
+  def self.find_or_create_versions(keys)
+    need_to_bump_versions = false
+
+    # Must check the local pending_memo_versions first in order to generate
+    # memo checksums. The pending_memo_versions are the expected versions that
+    # would be used if a transaction commited. With checksums consistent of
+    # pending versions, the method results would only be visible after a
+    # transaction commited (we bump the pending_memo_versions on redis as an
+    # after_commit callback)
+    if RedisMemo::AfterCommit.in_transaction?
+      memo_versions = RedisMemo::AfterCommit.pending_memo_versions.slice(*keys)
+    else
+      memo_versions = {}
+    end
+
+    keys_to_fetch = keys
+    keys_to_fetch -= memo_versions.keys unless memo_versions.empty?
+
+    cached_versions =
+      if keys_to_fetch.empty?
+        {}
+      else
+        RedisMemo::Cache.read_multi(
+          *keys_to_fetch,
+          raw: true,
+          raise_error: true,
+        )
+      end
+    memo_versions.merge!(cached_versions) unless cached_versions.empty?
+
+    versions = keys.map do |key|
+      version = memo_versions[key]
+      if version.nil?
+        # If a version does not exist, we assume it's because the version has
+        # expired due to TTL or it's evicted by a cache eviction policy. In
+        # this case, we will create a new version and use it for memoizing the
+        # cached result.
+        need_to_bump_versions = true
+
+        new_version = RedisMemo.uuid
+        RedisMemo::Memoizable::Invalidation.bump_version_later(
+          key,
+          new_version,
+          previous_version: '',
+        )
+        new_version
+      else
+        version
+      end
+    end
+
+    # Flush out the versions to Redis (async) if we created new versions
+    RedisMemo::Memoizable::Invalidation.drain_invalidation_queue if need_to_bump_versions
+
+    versions
+  end
+end

data/lib/redis_memo/memoizable/dependency.rb

@@ -0,0 +1,66 @@
+# frozen_string_literal: true
+
+# A Directed Acyclic Graph (DAG) of Memoizables
+class RedisMemo::Memoizable::Dependency
+  attr_accessor :nodes
+
+  def initialize
+    @nodes = {}
+  end
+
+  def memos
+    @nodes.values
+  end
+
+  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
+      [
+        dependency.redis_memo_class_memoizable,
+        RedisMemo::MemoizeQuery.create_memo(dependency, **conditions),
+      ].each do |memo|
+        nodes[memo.cache_key] = memo
+      end
+    else
+      raise(
+        RedisMemo::ArgumentError,
+        "Invalid dependency #{dependency}"
+      )
+    end
+  end
+
+  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
+
+  class UsingActiveRecord
+    def self.===(dependency)
+      RedisMemo::MemoizeQuery.using_active_record?(dependency)
+    end
+  end
+end

data/lib/redis_memo/memoizable/invalidation.rb

@@ -0,0 +1,142 @@
+# frozen_string_literal: true
+require_relative '../after_commit'
+require_relative '../cache'
+
+module RedisMemo::Memoizable::Invalidation
+  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
+  # the actual and expected previous_version on Redis match, to ensure eventual
+  # consistency: If the versions mismatch, we will use a new version that has
+  # not been associated with any cached_results.
+  #
+  #   - No invalid cached results will be read
+  #
+  #   - New memoized calculations will write back the fresh_results using the
+  #     new version as part of their checksums.
+  #
+  # Note: Cached data is not guaranteed to be consistent by design. Between the
+  # moment we should invalidate a version and the moment we actually
+  # invalidated a version, we would serve out-dated cached results, as if the
+  # operations that triggered the invalidation has not yet happened.
+  @@invalidation_queue = Queue.new
+
+  def self.bump_version_later(key, version, previous_version: nil)
+    if RedisMemo::AfterCommit.in_transaction?
+      previous_version ||= RedisMemo::AfterCommit.pending_memo_versions[key]
+    end
+
+    local_cache = RedisMemo::Cache.local_cache
+    if previous_version.nil? && local_cache&.include?(key)
+      previous_version = local_cache[key]
+    elsif RedisMemo::AfterCommit.in_transaction?
+      # 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,
+        raw: true,
+      )[key]
+    end
+
+    local_cache&.send(:[]=, key, version)
+    if RedisMemo::AfterCommit.in_transaction?
+      RedisMemo::AfterCommit.bump_memo_version_after_commit(
+        key,
+        version,
+        previous_version: previous_version,
+      )
+    else
+      @@invalidation_queue << Task.new(key, version, previous_version)
+    end
+  end
+
+  def self.drain_invalidation_queue
+    async_handler = RedisMemo::DefaultOptions.async
+    if async_handler.nil?
+      drain_invalidation_queue_now
+    else
+      async_handler.call do
+        drain_invalidation_queue_now
+      end
+    end
+  end
+
+  LUA_BUMP_VERSION = <<~LUA
+    local key = KEYS[1]
+    local expected_prev_version,
+          desired_new_version,
+          version_on_mismatch,
+          ttl = unpack(ARGV)
+
+    local actual_prev_version = redis.call('get', key)
+    local new_version = version_on_mismatch
+    local px = {}
+
+    if (not actual_prev_version and expected_prev_version == '') or expected_prev_version == actual_prev_version then
+      new_version = desired_new_version
+    end
+
+    if ttl ~= '' then
+      px = {'px', ttl}
+    end
+
+    return redis.call('set', key, new_version, unpack(px))
+  LUA
+
+  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: [task.key],
+        argv: [task.previous_version, task.version, RedisMemo.uuid, ttl],
+      )
+      RedisMemo::Tracer.set_tag(enqueue_to_finish: task.duration)
+    end
+  end
+
+  def self.drain_invalidation_queue_now
+    retry_queue = []
+    until @@invalidation_queue.empty?
+      task = @@invalidation_queue.pop
+      begin
+        bump_version(task)
+      rescue SignalException, Redis::BaseConnectionError,
+             ::ConnectionPool::TimeoutError
+        retry_queue << task
+      end
+    end
+  ensure
+    retry_queue.each { |task| @@invalidation_queue << task }
+  end
+
+  at_exit do
+    # The best effort
+    drain_invalidation_queue_now
+  end
+end

data/lib/redis_memo/memoize_method.rb

@@ -0,0 +1,133 @@
+# frozen_string_literal: true
+require_relative 'batch'
+require_relative 'future'
+require_relative 'memoizable'
+require_relative 'middleware'
+require_relative 'options'
+
+module RedisMemo::MemoizeMethod
+  def memoize_method(method_name, method_id: nil, **options, &depends_on)
+    method_name_without_memo = :"_redis_memo_#{method_name}_without_memo"
+    method_name_with_memo = :"_redis_memo_#{method_name}_with_memo"
+
+    alias_method method_name_without_memo, method_name
+
+    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
+        when NilClass
+          RedisMemo::MemoizeMethod.method_id(self, method_name)
+        when String, Symbol
+          method_id
+        else
+          method_id.call(self, *args)
+        end,
+        args,
+        dependent_memos,
+        options,
+        method_name_without_memo,
+      )
+
+      if RedisMemo::Batch.current
+        RedisMemo::Batch.current << future
+        return future
+      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)
+    is_class_method = ref.class == Class
+    class_name = is_class_method ? ref.name : ref.class.name
+
+    "#{class_name}#{is_class_method ? '::' : '#'}#{method_name}"
+  end
+
+  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)
+    dependency
+  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|
+      if memos[i]
+        method_cache_key_versions[i] = [method_id, memo_checksums[j]]
+        j += 1
+      else
+        ordered_method_args = method_args.map do |arg|
+          arg.is_a?(Hash) ? RedisMemo.deep_sort_hash(arg) : arg
+        end
+
+        method_cache_key_versions[i] = [
+          method_id,
+          RedisMemo.checksum(ordered_method_args.to_json),
+        ]
+      end
+    end
+
+    method_cache_key_versions.map do |method_id, method_cache_key_version|
+      # Example:
+      #
+      #   RedisMemo:MyModel#slow_calculation:<global cache version>:<local
+      #   cache version>
+      #
+      [
+        RedisMemo.name,
+        method_id,
+        RedisMemo::DefaultOptions.global_cache_key_version,
+        method_cache_key_version,
+      ].join(':')
+    end
+  rescue RedisMemo::Cache::Rescuable
+    nil
+  end
+end