Rake task to cleanup expired ActiveSession lookup keys

In some cases ActiveSession.cleanup was not called after authentication,
so for some user ActiveSession lookup keys grew without ever cleaning
up. This Rake task manually iterates over the lookup keys and removes
ones without existing ActiveSession.

app/models/active_session.rb

@@ -93,12 +93,12 @@ class ActiveSession
   end
   def self.list_sessions(user)
-    sessions_from_ids(session_ids_for_user(user))
+    sessions_from_ids(session_ids_for_user(user.id))
   end
-  def self.session_ids_for_user(user)
+  def self.session_ids_for_user(user_id)
     Gitlab::Redis::SharedState.with do |redis|
-      redis.smembers(lookup_key_name(user.id))
+      redis.smembers(lookup_key_name(user_id))
     end
   end
@@ -129,7 +129,7 @@ class ActiveSession
   end
   def self.cleaned_up_lookup_entries(redis, user)
-    session_ids = session_ids_for_user(user)
+    session_ids = session_ids_for_user(user.id)
     entries = raw_active_session_entries(session_ids, user.id)
     # remove expired keys.

changelogs/unreleased/64257-active_session_lookup_key_cleanup.yml

+---
+title: Rake task to cleanup expired ActiveSession lookup keys
+merge_request: 30668
+author:
+type: performance

doc/raketasks/cleanup.md

@@ -137,3 +137,13 @@ level with `NICENESS`. Below are the valid levels, but consult
 - `1` or `Realtime`
 - `2` or `Best-effort` (default)
 - `3` or `Idle`
+
+## Remove expired ActiveSession lookup keys
+
+```
+# omnibus-gitlab
+sudo gitlab-rake gitlab:cleanup:sessions:active_sessions_lookup_keys
+
+# installation from source
+bundle exec rake gitlab:cleanup:sessions:active_sessions_lookup_keys RAILS_ENV=production
+```

lib/tasks/gitlab/cleanup.rake

@@ -127,6 +127,58 @@ namespace :gitlab do
       end
     end
+    namespace :sessions do
+      desc "GitLab | Cleanup | Sessions | Clean ActiveSession lookup keys"
+      task active_sessions_lookup_keys: :gitlab_environment do
+        session_key_pattern = "#{Gitlab::Redis::SharedState::USER_SESSIONS_LOOKUP_NAMESPACE}:*"
+        last_save_check = Time.at(0)
+        wait_time = 10.seconds
+        cursor = 0
+        total_users_scanned = 0
+
+        Gitlab::Redis::SharedState.with do |redis|
+          begin
+            cursor, keys = redis.scan(cursor, match: session_key_pattern)
+            total_users_scanned += keys.count
+
+            if last_save_check < Time.now - 1.second
+              while redis.info('persistence')['rdb_bgsave_in_progress'] == '1'
+                puts "BGSAVE in progress, waiting #{wait_time} seconds"
+                sleep(wait_time)
+              end
+              last_save_check = Time.now
+            end
+
+            keys.each do |key|
+              user_id = key.split(':').last
+
+              lookup_key_count = redis.scard(key)
+
+              session_ids = ActiveSession.session_ids_for_user(user_id)
+              entries = ActiveSession.raw_active_session_entries(session_ids, user_id)
+              session_ids_and_entries = session_ids.zip(entries)
+
+              inactive_session_ids = session_ids_and_entries.map do |session_id, session|
+                session_id if session.nil?
+              end.compact
+
+              redis.pipelined do |conn|
+                inactive_session_ids.each do |session_id|
+                  conn.srem(key, session_id)
+                end
+              end
+
+              if inactive_session_ids
+                puts "deleted #{inactive_session_ids.count} out of #{lookup_key_count} lookup keys for User ##{user_id}"
+              end
+            end
+          end while cursor.to_i != 0
+
+          puts "--- All done! Total number of scanned users: #{total_users_scanned}"
+        end
+      end
+    end
+
     def remove?
       ENV['REMOVE'] == 'true'
     end

spec/models/active_session_spec.rb

@@ -114,7 +114,7 @@ RSpec.describe ActiveSession, :clean_gitlab_redis_shared_state do
         redis.sadd("session:lookup:user:gitlab:#{user.id}", session_ids)
       end
-      expect(ActiveSession.session_ids_for_user(user)).to eq(session_ids)
+      expect(ActiveSession.session_ids_for_user(user.id)).to eq(session_ids)
     end
   end

spec/tasks/gitlab/cleanup_rake_spec.rb

@@ -185,4 +185,34 @@ describe 'gitlab:cleanup rake tasks' do
       end
     end
   end
+
+  context 'sessions' do
+    describe 'gitlab:cleanup:sessions:active_sessions_lookup_keys', :clean_gitlab_redis_shared_state do
+      subject(:rake_task) { run_rake_task('gitlab:cleanup:sessions:active_sessions_lookup_keys') }
+
+      let!(:user) { create(:user) }
+      let(:existing_session_id) { '5' }
+
+      before do
+        Gitlab::Redis::SharedState.with do |redis|
+          redis.set("session:user:gitlab:#{user.id}:#{existing_session_id}",
+                    Marshal.dump(true))
+          redis.sadd("session:lookup:user:gitlab:#{user.id}", (1..10).to_a)
+        end
+      end
+
+      it 'runs the task without errors' do
+        expect { rake_task }.not_to raise_error
+      end
+
+      it 'removes expired active session lookup keys' do
+        Gitlab::Redis::SharedState.with do |redis|
+          lookup_key = "session:lookup:user:gitlab:#{user.id}"
+          expect { subject }.to change { redis.scard(lookup_key) }.from(10).to(1)
+          expect(redis.smembers("session:lookup:user:gitlab:#{user.id}")).to(
+            eql([existing_session_id]))
+        end
+      end
+    end
+  end
 end