Fix DtdUtils static field retaining disposed Project (#8658)

#8657

Author: jwren

CHANGELOG.md

@@ -1,3 +1,15 @@
+## Next
+
+### Added
+
+### Changed
+
+### Removed
+
+### Fixed
+
+- DtdUtils static field retaining disposed Project (#8658)
+
 ## 88.2.0
 
 ### Added

src/io/flutter/dart/DtdUtils.java

@@ -14,18 +14,18 @@
 import java.util.concurrent.*;
 
 public class DtdUtils {
-  private static final Map<Project, CompletableFuture<DartToolingDaemonService>> WAITERS = new ConcurrentHashMap<>();
+  // Visible for testing
+  static final Map<Project, CompletableFuture<DartToolingDaemonService>> WAITERS = new ConcurrentHashMap<>();
 
   public @NotNull CompletableFuture<DartToolingDaemonService> readyDtdService(@NotNull Project project) {
+    final DartToolingDaemonService dtdService = DartToolingDaemonService.getInstance(project);
+    if (dtdService.getWebSocketReady()) {
+      return CompletableFuture.completedFuture(dtdService);
+    }
+
     return WAITERS.computeIfAbsent(project, p -> {
-      final DartToolingDaemonService dtdService = DartToolingDaemonService.getInstance(project);
       CompletableFuture<DartToolingDaemonService> readyService = new CompletableFuture<>();
 
-      if (dtdService.getWebSocketReady()) {
-        readyService.complete(dtdService);
-        return readyService;
-      }
-
       final ScheduledExecutorService scheduler = AppExecutorUtil.getAppScheduledExecutorService();
 
       final ScheduledFuture<?> poll = scheduler.scheduleWithFixedDelay(() -> {
@@ -42,9 +42,17 @@ public class DtdUtils {
       readyService.whenComplete((s, t) -> {
         poll.cancel(false);
         timeout.cancel(false);
-        if (t != null) {
-          WAITERS.remove(p);
-        }
+        // Remove from waiters when done.
+        // We use the scheduler to ensure this runs after computeIfAbsent returns,
+        // in case readyService was completed synchronously inside computeIfAbsent.
+        // Although with the check above, synchronous completion here is unlikely,
+        // it's safer to be async or just rely on the fact that polling completes it async.
+        // If polling completes it, we are on a different thread, so computeIfAbsent is definitely done.
+        // If timeout completes it, we are on a different thread.
+        // So direct removal is safe IF completion happens async.
+        // The only sync completion risk is if we checked again inside and completed.
+        // But we don't check inside synchronously anymore (only in poll).
+        WAITERS.remove(project);
       });
 
       return readyService;

testSrc/unit/io/flutter/dart/DtdUtilsTest.java

@@ -0,0 +1,79 @@
+package io.flutter.dart;
+
+import com.intellij.openapi.project.Project;
+import com.intellij.testFramework.ServiceContainerUtil;
+import com.jetbrains.lang.dart.ide.toolingDaemon.DartToolingDaemonService;
+import io.flutter.testing.CodeInsightProjectFixture;
+import io.flutter.testing.Testing;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+import java.util.concurrent.CompletableFuture;
+
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.when;
+
+public class DtdUtilsTest {
+  @Rule
+  public final CodeInsightProjectFixture fixture = Testing.makeCodeInsightModule();
+
+  @Mock
+  private DartToolingDaemonService mockDtdService;
+
+  private Project project;
+
+  @Before
+  public void setUp() throws Exception {
+    MockitoAnnotations.initMocks(this);
+    project = fixture.getProject();
+    
+    // Register mock service
+    // We need to use the disposable from the fixture to ensure it gets cleaned up
+    ServiceContainerUtil.registerServiceInstance(project, DartToolingDaemonService.class, mockDtdService);
+  }
+
+  @Test
+  public void testReadyDtdServiceRemovesProjectFromWaiters() throws Exception {
+    when(mockDtdService.getWebSocketReady()).thenReturn(true);
+    
+    DtdUtils dtdUtils = new DtdUtils();
+    CompletableFuture<DartToolingDaemonService> future = dtdUtils.readyDtdService(project);
+    
+    assertTrue(future.isDone());
+    // Verify WAITERS map is empty
+    assertTrue("WAITERS map should be empty after service is ready", DtdUtils.WAITERS.isEmpty());
+  }
+
+  @Test
+  public void testWaitersMapIsClearedAfterAsyncCompletion() throws Exception {
+    when(mockDtdService.getWebSocketReady()).thenReturn(false);
+
+    DtdUtils dtdUtils = new DtdUtils();
+    CompletableFuture<DartToolingDaemonService> future = dtdUtils.readyDtdService(project);
+
+    // It should be in the map now
+    assertTrue("Project should be in WAITERS map", DtdUtils.WAITERS.containsKey(project));
+
+    // Simulate service becoming ready
+    when(mockDtdService.getWebSocketReady()).thenReturn(true);
+
+    // Wait for the poller to pick it up (it runs every 500ms)
+    // We can just wait on the future with a timeout
+    future.get(5, java.util.concurrent.TimeUnit.SECONDS);
+
+    assertTrue(future.isDone());
+
+    // Wait for the removal to happen (it happens in whenComplete, which might be slightly delayed if async)
+    // We can poll the map
+    long start = System.currentTimeMillis();
+    while (!DtdUtils.WAITERS.isEmpty() && System.currentTimeMillis() - start < 2000) {
+      Thread.sleep(100);
+    }
+
+    assertTrue("WAITERS map should be empty after async completion", DtdUtils.WAITERS.isEmpty());
+  }
+}