Bulk update user todos count cache

app/models/todo.rb

@@ -152,6 +152,20 @@ class Todo < ApplicationRecord
     def pluck_user_id
       pluck(:user_id)
     end
+
+    # Count todos grouped by user_id and state, using an UNION query
+    # so we can utilize the partial indexes for each state.
+    def count_grouped_by_user_id_and_state
+      grouped_count = select(:user_id, 'count(id) AS count').group(:user_id)
+
+      done = grouped_count.where(state: :done).select("'done' AS state")
+      pending = grouped_count.where(state: :pending).select("'pending' AS state")
+      union = unscoped.from_union([done, pending], remove_duplicates: false)
+
+      connection.select_all(union).each_with_object({}) do |row, counts|
+        counts[[row['user_id'], row['state']]] = row['count']
+      end
+    end
   end
 
   def resource_parent

app/services/todo_service.rb

@@ -47,7 +47,7 @@ class TodoService
 
     yield target
 
-    todo_users.each(&:update_todos_count_cache)
+    Users::UpdateTodoCountCacheService.new(todo_users).execute if todo_users.present?
   end
 
   # When we reassign an assignable object (issuable, alert) we should:
@@ -227,14 +227,16 @@ class TodoService
     users_with_pending_todos = pending_todos(users, attributes).pluck_user_id
     users.reject! { |user| users_with_pending_todos.include?(user.id) && Feature.disabled?(:multiple_todos, user) }
 
-    users.map do |user|
+    todos = users.map do |user|
       issue_type = attributes.delete(:issue_type)
       track_todo_creation(user, issue_type)
 
-      todo = Todo.create(attributes.merge(user_id: user.id))
-      user.update_todos_count_cache
-      todo
+      Todo.create(attributes.merge(user_id: user.id))
     end
+
+    Users::UpdateTodoCountCacheService.new(users).execute
+
+    todos
   end
 
   def new_issuable(issuable, author)

app/services/users/update_todo_count_cache_service.rb

+# frozen_string_literal: true
+
+module Users
+  class UpdateTodoCountCacheService < BaseService
+    QUERY_BATCH_SIZE = 10
+
+    attr_reader :users
+
+    # users - An array of User objects
+    def initialize(users)
+      @users = users
+    end
+
+    def execute
+      users.each_slice(QUERY_BATCH_SIZE) do |users_batch|
+        todo_counts = Todo.for_user(users_batch).count_grouped_by_user_id_and_state
+
+        users_batch.each do |user|
+          update_count_cache(user, todo_counts, :done)
+          update_count_cache(user, todo_counts, :pending)
+        end
+      end
+    end
+
+    private
+
+    def update_count_cache(user, todo_counts, state)
+      count = todo_counts.fetch([user.id, state.to_s], 0)
+      expiration_time = user.count_cache_validity_period
+
+      Rails.cache.write(['users', user.id, "todos_#{state}_count"], count, expires_in: expiration_time)
+    end
+  end
+end

changelogs/unreleased/bulk-update-user-todos-count-cache.yml

+---
+title: Avoid N+1 query when updating todo count cache
+merge_request: 57622
+author:
+type: performance

spec/models/todo_spec.rb

@@ -376,6 +376,22 @@ RSpec.describe Todo do
     end
   end
 
+  describe '.group_by_user_id_and_state' do
+    let_it_be(:user1) { create(:user) }
+    let_it_be(:user2) { create(:user) }
+
+    before do
+      create(:todo, user: user1, state: :pending)
+      create(:todo, user: user1, state: :pending)
+      create(:todo, user: user1, state: :done)
+      create(:todo, user: user2, state: :pending)
+    end
+
+    specify do
+      expect(Todo.count_grouped_by_user_id_and_state).to eq({ [user1.id, "done"] => 1, [user1.id, "pending"] => 2, [user2.id, "pending"] => 1 })
+    end
+  end
+
   describe '.any_for_target?' do
     it 'returns true if there are todos for a given target' do
       todo = create(:todo)

spec/services/todo_service_spec.rb

@@ -3,6 +3,8 @@
 require 'spec_helper'
 
 RSpec.describe TodoService do
+  include AfterNextHelpers
+
   let_it_be(:project) { create(:project, :repository) }
   let_it_be(:author) { create(:user) }
   let_it_be(:assignee) { create(:user) }
@@ -343,19 +345,19 @@ RSpec.describe TodoService do
 
     describe '#destroy_target' do
       it 'refreshes the todos count cache for users with todos on the target' do
-        create(:todo, target: issue, user: john_doe, author: john_doe, project: issue.project)
+        create(:todo, state: :pending, target: issue, user: john_doe, author: john_doe, project: issue.project)
 
-        expect_any_instance_of(User).to receive(:update_todos_count_cache).and_call_original
+        expect_next(Users::UpdateTodoCountCacheService, [john_doe]).to receive(:execute)
 
-        service.destroy_target(issue) { }
+        service.destroy_target(issue) { issue.destroy! }
       end
 
       it 'does not refresh the todos count cache for users with only done todos on the target' do
         create(:todo, :done, target: issue, user: john_doe, author: john_doe, project: issue.project)
 
-        expect_any_instance_of(User).not_to receive(:update_todos_count_cache)
+        expect(Users::UpdateTodoCountCacheService).not_to receive(:new)
 
-        service.destroy_target(issue) { }
+        service.destroy_target(issue) { issue.destroy! }
       end
 
       it 'yields the target to the caller' do
@@ -1099,13 +1101,9 @@ RSpec.describe TodoService do
   it 'updates cached counts when a todo is created' do
     issue = create(:issue, project: project, assignees: [john_doe], author: author)
 
-    expect(john_doe.todos_pending_count).to eq(0)
-    expect(john_doe).to receive(:update_todos_count_cache).and_call_original
+    expect_next(Users::UpdateTodoCountCacheService, [john_doe]).to receive(:execute)
 
     service.new_issue(issue, author)
-
-    expect(Todo.where(user_id: john_doe.id, state: :pending).count).to eq 1
-    expect(john_doe.todos_pending_count).to eq(1)
   end
 
   shared_examples 'updating todos state' do |state, new_state, new_resolved_by = nil|

spec/services/users/update_todo_count_cache_service_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Users::UpdateTodoCountCacheService do
+  describe '#execute' do
+    let_it_be(:user1) { create(:user) }
+    let_it_be(:user2) { create(:user) }
+
+    let_it_be(:todo1) { create(:todo, user: user1, state: :done) }
+    let_it_be(:todo2) { create(:todo, user: user1, state: :done) }
+    let_it_be(:todo3) { create(:todo, user: user1, state: :pending) }
+    let_it_be(:todo4) { create(:todo, user: user2, state: :done) }
+    let_it_be(:todo5) { create(:todo, user: user2, state: :pending) }
+    let_it_be(:todo6) { create(:todo, user: user2, state: :pending) }
+
+    it 'updates the todos_counts for users', :use_clean_rails_memory_store_caching do
+      Rails.cache.write(['users', user1.id, 'todos_done_count'], 0)
+      Rails.cache.write(['users', user1.id, 'todos_pending_count'], 0)
+      Rails.cache.write(['users', user2.id, 'todos_done_count'], 0)
+      Rails.cache.write(['users', user2.id, 'todos_pending_count'], 0)
+
+      expect { described_class.new([user1, user2]).execute }
+        .to change(user1, :todos_done_count).from(0).to(2)
+        .and change(user1, :todos_pending_count).from(0).to(1)
+        .and change(user2, :todos_done_count).from(0).to(1)
+        .and change(user2, :todos_pending_count).from(0).to(2)
+
+      Todo.delete_all
+
+      expect { described_class.new([user1, user2]).execute }
+        .to change(user1, :todos_done_count).from(2).to(0)
+        .and change(user1, :todos_pending_count).from(1).to(0)
+        .and change(user2, :todos_done_count).from(1).to(0)
+        .and change(user2, :todos_pending_count).from(2).to(0)
+    end
+
+    it 'avoids N+1 queries' do
+      control_count = ActiveRecord::QueryRecorder.new { described_class.new([user1]).execute }.count
+
+      expect { described_class.new([user1, user2]).execute }.not_to exceed_query_limit(control_count)
+    end
+
+    it 'executes one query per batch of users' do
+      stub_const("#{described_class}::QUERY_BATCH_SIZE", 1)
+
+      expect(ActiveRecord::QueryRecorder.new { described_class.new([user1]).execute }.count).to eq(1)
+      expect(ActiveRecord::QueryRecorder.new { described_class.new([user1, user2]).execute }.count).to eq(2)
+    end
+
+    it 'sets the cache expire time to the users count_cache_validity_period' do
+      allow(user1).to receive(:count_cache_validity_period).and_return(1.minute)
+      allow(user2).to receive(:count_cache_validity_period).and_return(1.hour)
+
+      expect(Rails.cache).to receive(:write).with(['users', user1.id, anything], anything, expires_in: 1.minute).twice
+      expect(Rails.cache).to receive(:write).with(['users', user2.id, anything], anything, expires_in: 1.hour).twice
+
+      described_class.new([user1, user2]).execute
+    end
+  end
+end