Unify collection caching code

6 files changed, 30 insertions, 87 deletions

diff --git a/app/controllers/api/v1/accounts_controller.rb b/app/controllers/api/v1/accounts_controller.rb
index ffa8b04fb..9a356196c 100644
--- a/app/controllers/api/v1/accounts_controller.rb
+++ b/app/controllers/api/v1/accounts_controller.rb
@@ -48,7 +48,7 @@ class Api::V1::AccountsController < ApiController
 
   def statuses
     @statuses = @account.statuses.paginate_by_max_id(DEFAULT_STATUSES_LIMIT, params[:max_id], params[:since_id]).to_a
-    @statuses = cache(@statuses)
+    @statuses = cache_collection(@statuses, Status)
 
     set_maps(@statuses)
     set_counters_maps(@statuses)
@@ -111,23 +111,4 @@ class Api::V1::AccountsController < ApiController
     @followed_by = Account.followed_by_map([@account.id], current_user.account_id)
     @blocking    = Account.blocking_map([@account.id], current_user.account_id)
   end
-
-  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] }.compact
-  end
 end
diff --git a/app/controllers/api/v1/notifications_controller.rb b/app/controllers/api/v1/notifications_controller.rb
index b23d7570d..a24e0beb7 100644
--- a/app/controllers/api/v1/notifications_controller.rb
+++ b/app/controllers/api/v1/notifications_controller.rb
@@ -8,7 +8,7 @@ class Api::V1::NotificationsController < ApiController
 
   def index
     @notifications = Notification.where(account: current_account).paginate_by_max_id(20, params[:max_id], params[:since_id])
-    @notifications = cache(@notifications)
+    @notifications = cache_collection(@notifications, Notification)
     statuses       = @notifications.select { |n| !n.target_status.nil? }.map(&:target_status)
 
     set_maps(statuses)
@@ -20,25 +20,4 @@ class Api::V1::NotificationsController < ApiController
 
     set_pagination_headers(next_path, prev_path)
   end
-
-  private
-
-  def cache(raw)
-    uncached_ids           = []
-    cached_keys_with_value = Rails.cache.read_multi(*raw.map(&:cache_key))
-
-    raw.each do |notification|
-      uncached_ids << notification.id unless cached_keys_with_value.key?(notification.cache_key)
-    end
-
-    unless uncached_ids.empty?
-      uncached = Notification.where(id: uncached_ids).with_includes.map { |n| [n.id, n] }.to_h
-
-      uncached.values.each do |notification|
-        Rails.cache.write(notification.cache_key, notification)
-      end
-    end
-
-    raw.map { |notification| cached_keys_with_value[notification.cache_key] || uncached[notification.id] }.compact
-  end
 end
diff --git a/app/controllers/api/v1/timelines_controller.rb b/app/controllers/api/v1/timelines_controller.rb
index 3debbdfc4..89e54e2cf 100644
--- a/app/controllers/api/v1/timelines_controller.rb
+++ b/app/controllers/api/v1/timelines_controller.rb
@@ -8,6 +8,7 @@ class Api::V1::TimelinesController < ApiController
 
   def home
     @statuses = Feed.new(:home, current_account).get(DEFAULT_STATUSES_LIMIT, params[:max_id], params[:since_id]).to_a
+    @statuses = cache_collection(@statuses)
 
     set_maps(@statuses)
     set_counters_maps(@statuses)
@@ -23,6 +24,7 @@ class Api::V1::TimelinesController < ApiController
 
   def mentions
     @statuses = Feed.new(:mentions, current_account).get(DEFAULT_STATUSES_LIMIT, params[:max_id], params[:since_id]).to_a
+    @statuses = cache_collection(@statuses)
 
     set_maps(@statuses)
     set_counters_maps(@statuses)
@@ -38,7 +40,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)
+    @statuses = cache_collection(@statuses)
 
     set_maps(@statuses)
     set_counters_maps(@statuses)
@@ -55,7 +57,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)
+    @statuses = cache_collection(@statuses)
 
     set_maps(@statuses)
     set_counters_maps(@statuses)
@@ -71,22 +73,7 @@ class Api::V1::TimelinesController < ApiController
 
   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] }.compact
+  def cache_collection(raw)
+    super(raw, Status)
   end
 end
diff --git a/app/controllers/application_controller.rb b/app/controllers/application_controller.rb
index b03a2cdea..ba0098c71 100644
--- a/app/controllers/application_controller.rb
+++ b/app/controllers/application_controller.rb
@@ -52,4 +52,23 @@ class ApplicationController < ActionController::Base
   def current_account
     @current_account ||= current_user.try(:account)
   end
+
+  def cache_collection(raw, klass)
+    uncached_ids           = []
+    cached_keys_with_value = Rails.cache.read_multi(*raw.map(&:cache_key))
+
+    raw.each do |item|
+      uncached_ids << item.id unless cached_keys_with_value.key?(item.cache_key)
+    end
+
+    unless uncached_ids.empty?
+      uncached = klass.where(id: uncached_ids).with_includes.map { |item| [item.id, item] }.to_h
+
+      uncached.values.each do |item|
+        Rails.cache.write(item.cache_key, item)
+      end
+    end
+
+    raw.map { |item| cached_keys_with_value[item.cache_key] || uncached[item.id] }.compact
+  end
 end
diff --git a/app/models/feed.rb b/app/models/feed.rb
index 45cb923d1..7b181d529 100644
--- a/app/models/feed.rb
+++ b/app/models/feed.rb
@@ -16,8 +16,8 @@ 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 = cache(unhydrated)
-      @statuses = unhydrated.map { |id| status_map[id] }.compact
+      status_map = Status.where(id: unhydrated).map { |s| [s.id, s] }.to_h
+      @statuses  = unhydrated.map { |id| status_map[id] }.compact
     end
 
     @statuses
@@ -25,29 +25,6 @@ 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/spec/controllers/admin/pubsubhubbub_controller_spec.rb b/spec/controllers/admin/pubsubhubbub_controller_spec.rb
index 4e8016b70..068bd09a6 100644
--- a/spec/controllers/admin/pubsubhubbub_controller_spec.rb
+++ b/spec/controllers/admin/pubsubhubbub_controller_spec.rb
@@ -4,7 +4,7 @@ require 'rails_helper'
 RSpec.describe Admin::PubsubhubbubController, type: :controller do
   describe 'GET #index' do
     before do
-      sign_in :user, Fabricate(:user, admin: true)
+      sign_in Fabricate(:user, admin: true), scope: :user
     end
 
     it 'returns http success' do