More test fixes: WeakThreadKeyMap concurrency fix and leak‑check test enhancements

Rails のテストスイートで発生していた接続リーク判定のフレークを根本的に解消し、Ruby 3.3.5 未満で使用される WeakThreadKeyMap の競合状態を修正しました。さらに、メンテナンス中の接続保持を正しく扱うテストケースを追加しています。

背景

WeakThreadKeyMap は、Ruby 3.3.5 より前のバージョンで ObjectSpace::WeakKeyMap が GC クラッシュを起こすために採用された代替実装です。しかし、スレッドが死亡した際のエントリ除去が select! に依存していたため、非同期エグゼキュータが存在するマルチスレッド環境でレースコンディションが発生していました。

先行 PR #57744 と #57750 では、メンテナンススレッドのチェックアウト順序を調整してフレークの抑止を試みましたが、別のテストで新たな不安定要因が確認されました。これにより、データ構造自体のスレッド安全性を高める根本的な対策が必要となりました。

この背景を受けて、WeakThreadKeyMap の内部実装をスレッド安全なマップに置き換え、同時実行時の競合を防止することが本 PR の主目的です。

技術的な変更

WeakThreadKeyMap の内部ストレージを Concurrent::Map に変更し、スレッド間で安全に参照・更新できるようにしました。これにより、select! による一括クリーニングから、個別にスレッドの存命確認を行うロジックへと移行します。

@@ -136,7 +136,7 @@ class WeakThreadKeyMap # :nodoc:
           # FIXME: On 3.3 we could use ObjectSpace::WeakKeyMap
           # but it currently causes GC crashes: https://github.com/byroot/rails/pull/3
           def initialize
-            @map = {}
+            @map = Concurrent::Map.new
           end

[]= メソッドでは、従来の @map.select! { |c, _| c&.alive? } を廃止し、each_pair で走査しながら死亡スレッドのエントリを削除するコードに置き換えました。変更後は明示的に delete を呼び出すため、同時実行時の不整合が起きません。

@@ -148,7 +148,9 @@ def []=(key, value)
-            @map.select! { |c, _| c&.alive? }
+            @map.each_pair do |thread, _|
+              @map.delete(thread) unless thread&.alive?
+            end
             @map[key] = value
           end

テスト側では check_connection_leaks にオプション引数 connection_pools = nil を導入し、必要に応じて対象プールだけを走査できるようにしました。既存の呼び出しは引数が省略されるため従来通り全プールをチェックします。

@@ -39,13 +39,13 @@ def after_teardown
       check_connection_leaks
     end
-
-    def check_connection_leaks
+    def check_connection_leaks(connection_pools = nil)
       return if in_memory_db?
@@
-      ActiveRecord::Base.connection_handler.each_connection_pool do |pool|
+      (connection_pools || ActiveRecord::Base.connection_handler.each_connection_pool).each do |pool|
         # Ensure all in flights tasks are completed.
         # Otherwise they may still hold a connection.
         if pool.async_executor

さらに、activerecord/test/cases/test_case_test.rb に2つのテストケースが追加され、メンテナンス中や非同期エグゼキュータが稼働中でも接続リークと誤判定しないことを検証しています。テストは Concurrent::Event でスレッド同期を制御し、レースコンディションを再現しつつ期待通りに動作することを確認します。

@@
-      release_maintenance_latch = Concurrent::CountDownLatch.new(1)
-      leak_check_attempted_latch = Concurrent::CountDownLatch.new(1)
+      release_maintenance = Concurrent::Event.new
+      leak_check_attempted = Concurrent::Event.new
@@
-      maintenance_thread = Thread.new do
-        Thread.current.name = "AR Pool Reaper"
-
-        pool.reaper_lock do
-          pool.send(:checkout_for_maintenance, conn)
-          release_maintenance_latch.wait
-          pool.send(:return_from_maintenance, conn)
-        end
-      end
+      maintenance_thread = Thread.new do
+        Thread.current.name = "AR Pool Reaper"
+
+        pool.reaper_lock do
+          pool.send(:checkout_for_maintenance, conn)
+          release_maintenance.wait
+          pool.send(:return_from_maintenance, conn)
+        end
+      end

これらの変更は、内部データ構造の安全性向上とテストの信頼性強化を同時に実現し、既存のアプリケーションコードへの影響はありません。

設計判断

Concurrent::Map への置き換えは、外部インターフェイス（[] と []=）を保ちつつスレッド安全性を確保する最小侵入的な選択でした。既存コードは API が変わらないため、利用側の改修は不要です。

check_connection_leaks に引数を追加したのは、テストから特定プール集合を注入して非同期エグゼキュータの影響を切り離すための設計です。デフォルトは nil で従来通り全プールを走査し、公開 API の互換性を維持しています。

新規テストでは Concurrent::Event や Concurrent::CountDownLatch といった同期プリミティブを用いて、レース条件を deterministic に再現できるようにしました。これにより、将来的に他の非同期機構が追加された際にも同様の検証が容易になります。

まとめ

WeakThreadKeyMap の内部実装を Concurrent::Map に置き換えることで、Ruby 3.3.5 未満での同時実行バグを根本的に解消しました。加えてテスト側の拡張と新規テストケースにより、接続リーク判定の誤検出が防止され、テストスイートの安定性が向上します。これらの変更は公開 API を破壊せず、既存アプリケーションへの影響を最小限に抑えた形で導入されています。