Implement includes caching for timelines APIs

3 files changed, 52 insertions, 9 deletions

diff --git a/app/controllers/api/v1/timelines_controller.rb b/app/controllers/api/v1/timelines_controller.rb
index c5ed315d9..b1d7c3052 100644
--- a/app/controllers/api/v1/timelines_controller.rb
+++ b/app/controllers/api/v1/timelines_controller.rb
@@ -38,6 +38,7 @@ class Api::V1::TimelinesController < ApiController
 
   def public
     @statuses = Status.as_public_timeline(current_account).paginate_by_max_id(DEFAULT_STATUSES_LIMIT, params[:max_id], params[:since_id]).to_a
+    @statuses = cache(@statuses)
 
     set_maps(@statuses)
     set_counters_maps(@statuses)
@@ -54,6 +55,7 @@ class Api::V1::TimelinesController < ApiController
   def tag
     @tag      = Tag.find_by(name: params[:id].downcase)
     @statuses = @tag.nil? ? [] : Status.as_tag_timeline(@tag, current_account).paginate_by_max_id(DEFAULT_STATUSES_LIMIT, params[:max_id], params[:since_id]).to_a
+    @statuses = cache(@statuses)
 
     set_maps(@statuses)
     set_counters_maps(@statuses)
@@ -66,4 +68,25 @@ class Api::V1::TimelinesController < ApiController
 
     render action: :index
   end
+
+  private
+
+  def cache(raw)
+    uncached_ids           = []
+    cached_keys_with_value = Rails.cache.read_multi(*raw.map(&:cache_key))
+
+    raw.each do |status|
+      uncached_ids << status.id unless cached_keys_with_value.key?(status.cache_key)
+    end
+
+    unless uncached_ids.empty?
+      uncached = Status.where(id: uncached_ids).with_includes.map { |s| [s.id, s] }.to_h
+
+      uncached.values.each do |status|
+        Rails.cache.write(status.cache_key, status)
+      end
+    end
+
+    raw.map { |status| cached_keys_with_value[status.cache_key] || uncached[status.id] }
+  end
 end
diff --git a/app/models/feed.rb b/app/models/feed.rb
index 6337726be..45cb923d1 100644
--- a/app/models/feed.rb
+++ b/app/models/feed.rb
@@ -16,7 +16,7 @@ class Feed
       RegenerationWorker.perform_async(@account.id, @type)
       @statuses = Status.send("as_#{@type}_timeline", @account).paginate_by_max_id(limit, nil, nil)
     else
-      status_map = Status.where(id: unhydrated).with_includes.map { |status| [status.id, status] }.to_h
+      status_map = cache(unhydrated)
       @statuses = unhydrated.map { |id| status_map[id] }.compact
     end
 
@@ -25,6 +25,29 @@ class Feed
 
   private
 
+  def cache(ids)
+    raw                    = Status.where(id: ids).to_a
+    uncached_ids           = []
+    cached_keys_with_value = Rails.cache.read_multi(*raw.map(&:cache_key))
+
+    raw.each do |status|
+      uncached_ids << status.id unless cached_keys_with_value.key?(status.cache_key)
+    end
+
+    unless uncached_ids.empty?
+      uncached = Status.where(id: uncached_ids).with_includes.map { |s| [s.id, s] }.to_h
+
+      uncached.values.each do |status|
+        Rails.cache.write(status.cache_key, status)
+      end
+    end
+
+    cached = cached_keys_with_value.values.map { |s| [s.id, s] }.to_h
+    cached.merge!(uncached) unless uncached_ids.empty?
+
+    cached
+  end
+
   def key
     FeedManager.instance.key(@type, @account.id)
   end
diff --git a/app/models/status.rb b/app/models/status.rb
index 708fef17f..3402929bf 100644
--- a/app/models/status.rb
+++ b/app/models/status.rb
@@ -27,7 +27,7 @@ class Status < ApplicationRecord
   default_scope { order('id desc') }
 
   scope :with_counters, -> { select('statuses.*, (select count(r.id) from statuses as r where r.reblog_of_id = statuses.id) as reblogs_count, (select count(f.id) from favourites as f where f.status_id = statuses.id) as favourites_count') }
-  scope :with_includes, -> { includes(:account, :media_attachments, :tags, :stream_entry, mentions: :account, reblog: [:account, :stream_entry, :tags, mentions: :account], thread: :account) }
+  scope :with_includes, -> { includes(:account, :media_attachments, :tags, :stream_entry, mentions: :account, reblog: [:account, :stream_entry, :tags, :media_attachments, mentions: :account], thread: :account) }
 
   def local?
     uri.nil?
@@ -89,28 +89,25 @@ class Status < ApplicationRecord
 
   class << self
     def as_home_timeline(account)
-      where(account: [account] + account.following).with_includes.with_counters
+      where(account: [account] + account.following).with_includes
     end
 
     def as_mentions_timeline(account)
-      where(id: Mention.where(account: account).pluck(:status_id)).with_includes.with_counters
+      where(id: Mention.where(account: account).pluck(:status_id)).with_includes
     end
 
     def as_public_timeline(account = nil)
       query = joins('LEFT OUTER JOIN accounts ON statuses.account_id = accounts.id').where('accounts.silenced = FALSE')
       query = filter_timeline(query, account) unless account.nil?
-
-      query.with_includes.with_counters
+      query
     end
 
     def as_tag_timeline(tag, account = nil)
       query = tag.statuses
                  .joins('LEFT OUTER JOIN accounts ON statuses.account_id = accounts.id')
                  .where('accounts.silenced = FALSE')
-
       query = filter_timeline(query, account) unless account.nil?
-
-      query.with_includes.with_counters
+      query
     end
 
     def favourites_map(status_ids, account_id)