From f7d6940e3a2168e76e3bc92a85bc6341bbfde021 Mon Sep 17 00:00:00 2001
From: Nazim Bhuiyan <nubhuiyan@ibm.com>
Date: Thu, 22 Feb 2024 13:32:36 -0500
Subject: [PATCH] Use Unsafe.compareAndSwap for writing to Stable field and
 array

Using Unsafe.compareAndSwap prevents issues resulting from multiple
threads coming across uninitialized VarHandle.typesAndInvokers
field and competing to populate the field with their own versions.
Using compareAndSwap to populate the typesAndInvokers'
MethodHandle table also prevents similar issues when multiple
threads come across the same unpopulated slot in the table and
compete to populate it.

Signed-off-by: Nazim Bhuiyan <nubhuiyan@ibm.com>
Co-authored-by: Devin Papineau <devin@ajdmp.ca>
---
 .../classes/java/lang/invoke/VarHandle.java   | 28 +++++++++++++++++--
 1 file changed, 26 insertions(+), 2 deletions(-)

diff --git a/src/java.base/share/classes/java/lang/invoke/VarHandle.java b/src/java.base/share/classes/java/lang/invoke/VarHandle.java
index 89d9e7f75d2..b5bc6a39f8d 100644
--- a/src/java.base/share/classes/java/lang/invoke/VarHandle.java
+++ b/src/java.base/share/classes/java/lang/invoke/VarHandle.java
@@ -23,6 +23,12 @@
  * questions.
  */
 
+/*
+ * ===========================================================================
+ * (c) Copyright IBM Corp. 2024, 2024 All Rights Reserved
+ * ===========================================================================
+ */
+
 package java.lang.invoke;
 
 import java.lang.constant.ClassDesc;
@@ -39,6 +45,7 @@
 import java.util.function.BiFunction;
 import java.util.function.Function;
 
+import jdk.internal.misc.Unsafe;
 import jdk.internal.util.Preconditions;
 import jdk.internal.vm.annotation.DontInline;
 import jdk.internal.vm.annotation.ForceInline;
@@ -2140,11 +2147,22 @@ static class TypesAndInvokers {
         MethodHandle[] methodHandle_table = new MethodHandle[AccessMode.COUNT];
     }
 
+    private static final long TYPES_AND_INVOKERS_OFFSET;
+
+    static {
+        TYPES_AND_INVOKERS_OFFSET = UNSAFE.objectFieldOffset(VarHandle.class, "typesAndInvokers");
+    }
+
     @ForceInline
     private final TypesAndInvokers getTypesAndInvokers() {
         TypesAndInvokers tis = typesAndInvokers;
         if (tis == null) {
-            tis = typesAndInvokers = new TypesAndInvokers();
+            tis = new TypesAndInvokers();
+            Object other = UNSAFE.compareAndExchangeReference(this, TYPES_AND_INVOKERS_OFFSET, null, tis);
+            if (other != null) {
+                // Lost the race, so use what was set by winning thread.
+                tis = (TypesAndInvokers) other;
+            }
         }
         return tis;
     }
@@ -2154,7 +2172,13 @@ MethodHandle getMethodHandle(int mode) {
         TypesAndInvokers tis = getTypesAndInvokers();
         MethodHandle mh = tis.methodHandle_table[mode];
         if (mh == null) {
-            mh = tis.methodHandle_table[mode] = getMethodHandleUncached(mode);
+            mh = getMethodHandleUncached(mode);
+            long offset = Unsafe.ARRAY_OBJECT_BASE_OFFSET + (Unsafe.ARRAY_OBJECT_INDEX_SCALE * mode);
+            Object other = UNSAFE.compareAndExchangeReference(tis.methodHandle_table, offset, null, mh);
+            if (other != null) {
+                // We lost the race. Use the winning thread's handle instead.
+                mh = (MethodHandle) other;
+            }
         }
         return mh;
     }