mongoid_timestamp_cacheable.rb

module Mongoid
  # Mongoid extension which stores all fetched records in a memory cache
  module TimestampCacheable
    extend ActiveSupport::Concern

    class QueryNotCacheableError < StandardError; end

    module ClassMethods
      def queryable
        scope_stack.last || TimestampCacheableCriteria.new(self)
      end

      def __timestamp_cache__
        @_cache ||= ActiveSupport::Cache::MemoryStore.new
      end
    end
  end

  class TimestampCacheableCriteria < Mongoid::Criteria
    def each(&_block)
      raise QueryNotCacheableError, 'Don\'t specify fields in your query for timestamp cacheable models' if query.operation.fields.present?

      prefetch = query.select(updated_at: 1, _id: 1).to_a
      refetched = refetch_stale_records(prefetch)

      prefetch.each do |id_doc|
        # Prefer cached document and fallback on refetched document
        doc = __timestamp_cache__.fetch(document_cache_key(id_doc)) ||
                refetched.detect { |d| d['_id'] == id_doc['_id'] }

        # Hydrate cached/refetched document attribute hash as model, recache, and yield
        model = build_model(doc)
        __timestamp_cache__.write(model.cache_key, doc)

        yield model if block_given?
      end
    end

    private

    def refetch_stale_records(prefetch)
      klass.collection.find('_id' => { '$in' => document_ids_to_refetch(prefetch) }).to_a
    end

    def document_ids_to_refetch(prefetch)
      result = []
      prefetch.each do |id_doc|
        cache_key = document_cache_key(id_doc)
        cached_doc = __timestamp_cache__.fetch(cache_key)
        if cached_doc.nil? || cached_doc['updated_at'] < id_doc['updated_at']
          __timestamp_cache__.delete(cache_key)
          result << id_doc['_id']
        end
      end

      result
    end

    def document_cache_key(doc)
      build_model(doc).cache_key
    end

    def build_model(doc)
      Mongoid::Factory.from_db(klass, doc)
    end
  end
end