From b62cd2383dcd1f4857fc4cdd269b257eef5cd432 Mon Sep 17 00:00:00 2001
From: Alexander Kuznetsov <aleksander.kuznetsov@jetbrains.com>
Date: Mon, 16 Jun 2025 18:37:27 +0200
Subject: [PATCH] IDEA-374618 [debugger] Attempt to fix race condition leading
 to AIOOBE: Int2ObjectOpenHashMap.get

GitOrigin-RevId: 9d9ed196a4e72922019f482a10d260d031be3680
---
 .../ui/impl/watch/MethodsTracker.java         | 34 +++++++++++++++++--
 .../ui/impl/watch/canDropFrameUtils.kt        |  4 +--
 2 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/java/debugger/impl/src/com/intellij/debugger/ui/impl/watch/MethodsTracker.java b/java/debugger/impl/src/com/intellij/debugger/ui/impl/watch/MethodsTracker.java
index 2cf38afdbe7b..a21cdc04d9eb 100644
--- a/java/debugger/impl/src/com/intellij/debugger/ui/impl/watch/MethodsTracker.java
+++ b/java/debugger/impl/src/com/intellij/debugger/ui/impl/watch/MethodsTracker.java
@@ -2,7 +2,6 @@
 package com.intellij.debugger.ui.impl.watch;
 
 import com.sun.jdi.Method;
-import it.unimi.dsi.fastutil.ints.Int2ObjectMap;
 import it.unimi.dsi.fastutil.ints.Int2ObjectOpenHashMap;
 import it.unimi.dsi.fastutil.objects.Object2IntOpenHashMap;
 import org.jetbrains.annotations.Nullable;
@@ -14,7 +13,7 @@ import java.util.concurrent.CompletableFuture;
  */
 public class MethodsTracker {
   @SuppressWarnings("SSBasedInspection") private final Object2IntOpenHashMap<Method> myMethodCounter = new Object2IntOpenHashMap<>();
-  private final Int2ObjectMap<MethodOccurrence> myCache = new Int2ObjectOpenHashMap<>();
+  private final MyCache<MethodOccurrence> myCache = new MyCache<>();
   private final CompletableFuture<Void> myFinished = new CompletableFuture<>();
 
   public final class MethodOccurrence {
@@ -52,9 +51,13 @@ public class MethodsTracker {
       return myMethod != null && getOccurrenceCount(myMethod) > 1;
     }
 
-    MethodOccurrence getMethodOccurrence(int frameIndex) {
+    MethodOccurrence getMethodOccurrenceSync(int frameIndex) {
       return myCache.get(frameIndex);
     }
+
+    CompletableFuture<MethodOccurrence> getMethodOccurrenceAsync(int frameIndex) {
+      return myCache.getExact(frameIndex);
+    }
   }
 
   public void finish() {
@@ -79,4 +82,29 @@ public class MethodsTracker {
   private CompletableFuture<Integer> getExactOccurrenceCount(@Nullable Method method) {
     return myFinished.thenApply(__ -> myMethodCounter.getInt(method));
   }
+
+  /**
+   * Quasi-thread-safe map that works on the following premises:
+   * <ul>
+   *   <li>Only {@code computeIfAbsent} and {@code get} are used</li>
+   *   <li>{@code computeIfAbsent} is called sequentially within the same thread (and so is {@code rehash})</li>
+   *   <li>{@code get} can be called in parallel, but it's guaranteed that the map
+   *  will eventually contain the bucket for the key.</li>
+   * </ul>
+   */
+  private static class MyCache<V> extends Int2ObjectOpenHashMap<V> {
+
+    private volatile CompletableFuture<Void> myRehashFinished = CompletableFuture.completedFuture(null);
+
+    public CompletableFuture<V> getExact(int key) {
+      return myRehashFinished.thenApply(__ -> super.get(key));
+    }
+
+    @Override
+    protected void rehash(int newN) {
+      myRehashFinished = new CompletableFuture<>();
+      super.rehash(newN);
+      myRehashFinished.complete(null);
+    }
+  }
 }
diff --git a/java/debugger/impl/src/com/intellij/debugger/ui/impl/watch/canDropFrameUtils.kt b/java/debugger/impl/src/com/intellij/debugger/ui/impl/watch/canDropFrameUtils.kt
index a59088c69a9d..2822bf88047d 100644
--- a/java/debugger/impl/src/com/intellij/debugger/ui/impl/watch/canDropFrameUtils.kt
+++ b/java/debugger/impl/src/com/intellij/debugger/ui/impl/watch/canDropFrameUtils.kt
@@ -15,7 +15,7 @@ import java.util.concurrent.CompletableFuture
 
 internal fun StackFrameDescriptorImpl.canDropFrameSync(): ThreeState {
   return isSafeToDropFrame(uiIndex, unsureIfCallerFrameAbsent = true) { i ->
-    methodOccurrence?.getMethodOccurrence(i)?.method
+    methodOccurrence?.getMethodOccurrenceSync(i)?.method
   }
 }
 
@@ -42,7 +42,7 @@ internal fun StackFrameDescriptorImpl.canDropFrameAsync(): CompletableFuture<Boo
     }
 
     isSafeToDropFrame(uiIndex, unsureIfCallerFrameAbsent = false) { i ->
-      methodOccurrence?.getMethodOccurrence(i)?.method ?: computeMethod(i)
+      methodOccurrence?.getMethodOccurrenceAsync(i)?.await()?.method ?: computeMethod(i)
     }.toBoolean()
   }
 }