8351045: ClassValue::remove cannot ensure computation observes up-to-date state

Reviewed-by: rriggs, jrose

src/java.base/share/classes/java/lang/ClassValue.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2010, 2024, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2010, 2025, Oracle and/or its affiliates. All rights reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * This code is free software; you can redistribute it and/or modify it
@@ -500,15 +500,15 @@ public abstract class ClassValue<T> {
         synchronized
         void removeEntry(ClassValue<?> classValue) {
             Entry<?> e = remove(classValue.identity);
-            if (e == null) {
-                // Uninitialized, and no pending calls to computeValue.  No change.
-            } else if (e.isPromise()) {
-                // State is uninitialized, with a pending call to finishEntry.
-                // Since remove is a no-op in such a state, keep the promise
-                // by putting it back into the map.
-                put(classValue.identity, e);
-            } else {
-                // In an initialized state.  Bump forward, and de-initialize.
+            // e == null: Uninitialized, and no pending calls to computeValue.
+            // remove(identity) didn't change anything.  No change.
+            // e.isPromise(): computeValue already used outdated values.
+            // remove(identity) discarded the outdated computation promise.
+            // finishEntry will retry when it discovers the promise is removed.
+            // No cache invalidation.  No further action needed.
+            if (e != null && !e.isPromise()) {
+                // Initialized.
+                // Bump forward to invalidate racy-read cached entries.
                 classValue.bumpVersion();
                 // Make all cache elements for this guy go stale.
                 removeStaleEntries(classValue);

test/jdk/java/lang/invoke/ClassValueTest.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2011, 2013, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2011, 2025, Oracle and/or its affiliates. All rights reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * This code is free software; you can redistribute it and/or modify it
@@ -21,22 +21,27 @@
  * questions.
  */
 
-/* @test
+/*
+ * @test
+ * @bug 8351045
  * @summary tests for class-specific values
- * @compile ClassValueTest.java
- * @run testng/othervm test.java.lang.invoke.ClassValueTest
+ * @run junit ClassValueTest
  */
 
-package test.java.lang.invoke;
+import java.util.ArrayList;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.atomic.AtomicInteger;
 
-import org.testng.*;
-import static org.testng.AssertJUnit.*;
-import org.testng.annotations.*;
+import org.junit.jupiter.api.Test;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
 
 /**
  * @author jrose
  */
-public class ClassValueTest {
+final class ClassValueTest {
     static String nameForCV1(Class<?> type) {
         return "CV1:" + type.getName();
     }
@@ -143,7 +148,7 @@ public class ClassValueTest {
                 }
             }
         }
-        assertEquals(countForCVN, 0);
+        assertEquals(0, countForCVN);
         System.out.println("[rechecking values]");
         for (int i = 0; i < cvns.length * 10; i++) {
             int n = i % cvns.length;
@@ -152,4 +157,57 @@ public class ClassValueTest {
             }
         }
     }
+
+    private static final int RUNS = 16;
+    private static final long COMPUTE_TIME_MILLIS = 100;
+
+    @Test
+    void testRemoveOnComputeCases() {
+        try (var exec = Executors.newVirtualThreadPerTaskExecutor()) {
+            var tasks = new ArrayList<Future<?>>(RUNS);
+            for (int i = 0; i < RUNS; i++) {
+                tasks.add(exec.submit(this::testRemoveOnCompute));
+            }
+            for (var task : tasks) {
+                try {
+                    task.get();
+                } catch (InterruptedException | ExecutionException ex) {
+                    var cause = ex.getCause();
+                    if (cause instanceof AssertionError ae)
+                        throw ae;
+                    throw new RuntimeException(ex);
+                }
+            }
+        }
+    }
+
+    void testRemoveOnCompute() {
+        AtomicInteger input = new AtomicInteger(0);
+        ClassValue<Integer> cv = new ClassValue<>() {
+            @Override
+            protected Integer computeValue(Class<?> type) {
+                // must get early to represent using outdated input
+                int v = input.get();
+                try {
+                    Thread.sleep(COMPUTE_TIME_MILLIS);
+                } catch (InterruptedException ex) {
+                    throw new RuntimeException(ex);
+                }
+                return v;
+            }
+        };
+        var innocuous = Thread.startVirtualThread(() -> cv.get(int.class));
+        var refreshInput = Thread.startVirtualThread(() -> {
+            input.incrementAndGet();
+            cv.remove(int.class); // Let's recompute with updated inputs!
+        });
+        try {
+            innocuous.join();
+            refreshInput.join();
+        } catch (InterruptedException ex) {
+            throw new RuntimeException(ex);
+        }
+        assertEquals(1, input.get(), "input not updated");
+        assertEquals(1, cv.get(int.class), "CV not using up-to-date input");
+    }
 }