Poor man's ActiveRecord identity map

identity_map.rb

# Poor man's identity map for reusing ActiveRecord objects and avoid fetching the same
# record multiple times from the database within a request or background job. The check method can
# be used to locate duplicate queries in development. The goal is to instantiate
# each ActiveRecord object only once. The Rails identity map is known to have issues:
# http://stackoverflow.com/questions/6905891/rails-3-1-identity-map-issues
class IdentityMap
  # Set one or more ActiveRecord objects in the map
  def set(*objects)
    objects.each do |object|
      map[key(object)] = object
    end
    objects.last
  end

  # Invoked with either just an ActiveRecord object, or an ActiveRecord class and id
  def get(object_or_class, object_id = nil)
    map[key(object_or_class, object_id)]
  end

  def fetch(object)
    return nil if object.nil?
    get(object) || set(object)
  end

  def contains?(object)
    get(object) && get(object).object_id == object.object_id
  end

  def check(object)
    return unless Rails.env.development? # only intended to be used in development
    unless contains?(object)
      Rails.logger.info("IdentityMap: duplicate object #{key(object)} #{object.object_id} (!= #{get(object).object_id}), from: #{caller[0, 4].inspect}")
    end
  end

  private

  def key(object_or_class, object_id = nil)
    if object_id
      klass = object_or_class
    else
      klass = object_or_class.class
      object_id = object_or_class.id
    end
    [klass.name, object_id].join(":")
  end

  def map
    @map ||= {}
  end
end