Fix timeouts on expiring SSH keys

Add keyset pagination

Changelog: fixed

app/models/key.rb

@@ -7,7 +7,7 @@ class Key < ApplicationRecord
   include Sortable
   include Sha256Attribute
   include Expirable
-  include EachBatch
+  include FromUnion
 
   sha256_attribute :fingerprint_sha256
 

app/workers/ssh_keys/expired_notification_worker.rb

@@ -11,12 +11,31 @@ module SshKeys
     tags :exclude_from_kubernetes
     idempotent!
 
+    BATCH_SIZE = 500
+
+    # rubocop: disable CodeReuse/ActiveRecord
     def perform
       return unless ::Feature.enabled?(:ssh_key_expiration_email_notification, default_enabled: :yaml)
 
-      # rubocop:disable CodeReuse/ActiveRecord
-      Key.expired_and_not_notified.each_batch(of: 1000) do |relation| # rubocop:disable Cop/InBatches
-        users = User.where(id: relation.select(:user_id))
+      order = Gitlab::Pagination::Keyset::Order.build([
+        Gitlab::Pagination::Keyset::ColumnOrderDefinition.new(
+          attribute_name: 'expires_at_utc',
+          order_expression: Arel.sql("date(expires_at AT TIME ZONE 'UTC')").asc,
+          nullable: :not_nullable,
+          distinct: false,
+          add_to_projections: true
+        ),
+        Gitlab::Pagination::Keyset::ColumnOrderDefinition.new(
+          attribute_name: 'id',
+          order_expression: Key.arel_table[:id].asc
+        )
+      ])
+
+      scope = Key.expired_and_not_notified.order(order)
+
+      iterator = Gitlab::Pagination::Keyset::Iterator.new(scope: scope, use_union_optimization: true)
+      iterator.each_batch(of: BATCH_SIZE) do |relation|
+        users = User.where(id: relation.map(&:user_id)) # Keyset pagination will load the rows
 
         users.each do |user|
           with_context(user: user) do
@@ -24,7 +43,7 @@ module SshKeys
           end
         end
       end
-      # rubocop:enable CodeReuse/ActiveRecord
     end
+    # rubocop: enable CodeReuse/ActiveRecord
   end
 end

db/migrate/20210609090856_add_expiry_id_ssh_key_notification_index.rb

+# frozen_string_literal: true
+
+class AddExpiryIdSshKeyNotificationIndex < ActiveRecord::Migration[6.1]
+  include Gitlab::Database::MigrationHelpers
+
+  disable_ddl_transaction!
+
+  INDEX_NAME = 'index_keys_on_expires_at_and_id'
+
+  def up
+    add_concurrent_index :keys,
+                         "date(timezone('UTC', expires_at)), id",
+                         where: 'expiry_notification_delivered_at IS NULL',
+                         name: INDEX_NAME
+  end
+
+  def down
+    remove_concurrent_index_by_name :keys, INDEX_NAME
+  end
+end

db/schema_migrations/20210609090856

+597e04c51bdad1900b2535c9d664c9e3a4d2a5879e657ef470bbc7ac461d3cca
\ No newline at end of file

db/structure.sql

@@ -23615,6 +23615,8 @@ CREATE INDEX index_jira_tracker_data_on_service_id ON jira_tracker_data USING bt
 
 CREATE INDEX index_keys_on_expires_at_and_expiry_notification_undelivered ON keys USING btree (date(timezone('UTC'::text, expires_at)), expiry_notification_delivered_at) WHERE (expiry_notification_delivered_at IS NULL);
 
+CREATE INDEX index_keys_on_expires_at_and_id ON keys USING btree (date(timezone('UTC'::text, expires_at)), id) WHERE (expiry_notification_delivered_at IS NULL);
+
 CREATE UNIQUE INDEX index_keys_on_fingerprint ON keys USING btree (fingerprint);
 
 CREATE INDEX index_keys_on_fingerprint_sha256 ON keys USING btree (fingerprint_sha256);

spec/workers/ssh_keys/expired_notification_worker_spec.rb

@@ -15,6 +15,20 @@ RSpec.describe SshKeys::ExpiredNotificationWorker, type: :worker do
   describe '#perform' do
     let_it_be(:user) { create(:user) }
 
+    context 'with a large batch' do
+      before do
+        stub_const("SshKeys::ExpiredNotificationWorker::BATCH_SIZE", 5)
+      end
+
+      let_it_be_with_reload(:keys) { create_list(:key, 20, expires_at: 3.days.ago, user: user) }
+
+      it 'updates all keys regardless of batch size' do
+        worker.perform
+
+        expect(keys.pluck(:expiry_notification_delivered_at)).not_to include(nil)
+      end
+    end
+
     context 'with expiring key today' do
       let_it_be_with_reload(:expired_today) { create(:key, expires_at: Time.current, user: user) }