chore: minimal complexity reduction in unsafe replacement

wrapper-jvm/impl/src/main/java/eu/cloudnetservice/wrapper/impl/transform/unsafe/ArrayOps.java

@@ -303,18 +303,19 @@ public static boolean arrayComparePut(
       // cant add to a boolean or object array, but can get from it
       var handle = arrayKind.arrayElementVarHandle();
       var arrayIndex = Math.toIntExact(offset / arrayKind.byteSize());
-      var witness = handle.compareAndExchange(array, arrayIndex, expected, value);
-      return kind.areValuesEqual(expected, witness);
+      return handle.compareAndSet(array, arrayIndex, expected, value);
     }
 
     // check if aligned access to the elements in the array is possible, that is
     // if an element can read from the array without requiring two operations to do so
     var arrayMemorySegment = arrayKind.createArrayMemorySegment(array);
     if (isAlignedAccess(arrayKind, kind, offset)) {
       var handle = kind.layoutVarHandle();
-      var witness = handle.compareAndExchange(arrayMemorySegment, offset, expected, value);
-      copyByteSegmentToBoolArray(arrayMemorySegment, array, offset, kind.byteSize());
-      return kind.areValuesEqual(expected, witness);
+      var result = handle.compareAndSet(arrayMemorySegment, offset, expected, value);
+      if (result) {
+        copyByteSegmentToBoolArray(arrayMemorySegment, array, offset, kind.byteSize());
+      }
+      return result;
     } else {
       var handle = kind.unalignedLayoutVarHandle();
       var current = handle.get(arrayMemorySegment, offset);

wrapper-jvm/impl/src/main/java/eu/cloudnetservice/wrapper/impl/transform/unsafe/FieldOps.java

@@ -316,10 +316,9 @@ static boolean fieldComparePut(
     } else {
       // field is non-final, can use var handles
       var handle = lookup.unreflectVarHandle(field);
-      var witness = isStatic
-        ? handle.compareAndExchange(convertedExp, convertedVal)
-        : handle.compareAndExchange(instance, convertedExp, convertedVal);
-      return kind.areValuesEqual(witness, convertedExp);
+      return isStatic
+        ? handle.compareAndSet(convertedExp, convertedVal)
+        : handle.compareAndSet(instance, convertedExp, convertedVal);
     }
   }
 }

wrapper-jvm/impl/src/main/java/eu/cloudnetservice/wrapper/impl/transform/unsafe/MemoryOps.java

@@ -172,8 +172,7 @@ public static boolean memComparePut(
   ) {
     if (isAlignedAccess(kind, offset)) {
       var handle = kind.layoutVarHandle();
-      var witness = handle.compareAndExchange(FULL_MEMORY_SEGMENT, offset, expected, value);
-      return kind.areValuesEqual(expected, witness);
+      return handle.compareAndSet(FULL_MEMORY_SEGMENT, offset, expected, value);
     } else {
       var handle = kind.unalignedLayoutVarHandle();
       var current = handle.get(FULL_MEMORY_SEGMENT, offset);

wrapper-jvm/impl/src/main/java/eu/cloudnetservice/wrapper/impl/transform/unsafe/UnsafeReplacementDelegate.java

@@ -29,10 +29,10 @@
 import java.util.Objects;
 import java.util.concurrent.locks.LockSupport;
 import java.util.function.Consumer;
-import java.util.function.Function;
 import java.util.function.Supplier;
 import lombok.NonNull;
 import org.jetbrains.annotations.Nullable;
+import org.jetbrains.annotations.VisibleForTesting;
 
 /**
  * Delegate class that holds current replacements for fields/methods defined in {@code sun.misc.Unsafe}.
@@ -46,9 +46,12 @@
 @Deprecated
 public final class UnsafeReplacementDelegate {
 
-  // accessors for cleaning up direct byte buffers
-  private static final Supplier<Consumer<ByteBuffer>> BB_CLEANER = createByteBufferCleaner();
-  private static final Supplier<Function<ByteBuffer, Object>> BB_ATTACHMENT_GETTER = createByteBufferAttachmentGetter();
+  // accessor for cleaning up direct byte buffers
+  @VisibleForTesting
+  static final Consumer<ByteBuffer> BB_CLEANER_NOOP = _ -> {
+  };
+  @VisibleForTesting
+  static final Supplier<Consumer<ByteBuffer>> BB_CLEANER = createByteBufferCleaner();
 
   // accessor for the operating system mx bean
   private static final Supplier<OperatingSystemMXBean> OS_MX_BEAN =
@@ -93,35 +96,6 @@ private UnsafeReplacementDelegate() {
     });
   }
 
-  /**
-   * Get a supplier that creates a function to resolve the attachment of a direct byte buffer. The function is only
-   * created once on the first initialization of the supplier.
-   *
-   * @return a supplier that creates a function to resolve the attachment of a direct byte buffer.
-   */
-  private static @NonNull Supplier<Function<ByteBuffer, Object>> createByteBufferAttachmentGetter() {
-    return new LazyMemoizingSupplier<>(() -> {
-      try {
-        var lookup = OpConstants.TRUSTED_LOOKUP.get();
-        var directBufferClass = Class.forName("sun.nio.ch.DirectBuffer");
-        var attachementMethod = directBufferClass.getMethod("attachment");
-        var attachmentMethodHandle = MethodHandles.explicitCastArguments(
-          lookup.unreflect(attachementMethod),
-          MethodType.methodType(Object.class, ByteBuffer.class));
-        return buffer -> {
-          try {
-            return attachmentMethodHandle.invokeExact(buffer);
-          } catch (Throwable _) {
-            return null;
-          }
-        };
-      } catch (Throwable throwable) {
-        UnsafeLogUtil.debug("Unable to access byte buffer attachment method; assuming it's never attached", throwable);
-        return _ -> null;
-      }
-    });
-  }
-
   /**
    * Get a supplier that creates a consumer to clean a direct byte buffer. The consumer is only created once on the
    * first initialization of the supplier.
@@ -133,31 +107,60 @@ private UnsafeReplacementDelegate() {
       try {
         var lookup = OpConstants.TRUSTED_LOOKUP.get();
 
-        // get the method handle to get the cleaner of the provided byte buffer (type: (ByteBuffer): Cleaner)
+        // get the method handle to get the cleaner of the provided byte buffer (type: (ByteBuffer):Cleaner)
         var directBufferClass = Class.forName("sun.nio.ch.DirectBuffer");
         var cleanerMethod = directBufferClass.getDeclaredMethod("cleaner");
         var cleanerHandle = MethodHandles.explicitCastArguments(
           lookup.unreflect(cleanerMethod),
           MethodType.methodType(cleanerMethod.getReturnType(), ByteBuffer.class));
 
-        // get the method handle to invoke the clean method on the Cleaner class (type: (Cleaner): void)
+        // get the method handle to invoke the clean method on the Cleaner class (type: (Cleaner):void)
         var cleanerClass = Class.forName("jdk.internal.ref.Cleaner");
         var cleanMethod = cleanerClass.getDeclaredMethod("clean");
         var cleanHandle = lookup.unreflect(cleanMethod);
 
-        // adapt the clean() method handle by pre-processing it with the result of the cleaner retrieval method handle
-        // this results in a chained invocation like clean(buffer.cleaner())
-        var cleanBufferHandle = MethodHandles.filterArguments(cleanHandle, 0, cleanerHandle);
+        // adapt the clean() method handle by pre-processing it with the result of the cleaner retrieval
+        // method handle, this results in a chained invocation (type: (ByteBuffer):void)
+        var cleanBufferHandle = MethodHandles.filterReturnValue(cleanerHandle, cleanHandle);
+
+        // get a method handle to check if a buffer has an attachment (type: (ByteBuffer):boolean)
+        var attachementMethod = directBufferClass.getMethod("attachment");
+        var attachmentHandle = MethodHandles.explicitCastArguments(
+          lookup.unreflect(attachementMethod),
+          MethodType.methodType(attachementMethod.getReturnType(), ByteBuffer.class));
+        var isNullHandle = lookup.findStatic(
+          Objects.class,
+          "isNull",
+          MethodType.methodType(boolean.class, Object.class));
+        var isAttachmentNullHandle = MethodHandles.collectArguments(isNullHandle, 0, attachmentHandle);
+
+        // get a method handle that throws an IAE with the message 'duplicate or slice' (type: ():void)
+        // this handle needs to then be adapted to add an extra, ignored ByteBuffer param (type: (ByteBuffer):void)
+        var iaeMessageCtrHandle = lookup.findConstructor(
+          IllegalArgumentException.class,
+          MethodType.methodType(void.class, String.class));
+        var dupOrSliceIaeCtrHandle = MethodHandles.insertArguments(iaeMessageCtrHandle, 0, "duplicate or slice");
+        var throwIaeHandle = MethodHandles.throwException(void.class, IllegalArgumentException.class);
+        var throwDupOrSliceHandle = MethodHandles.collectArguments(throwIaeHandle, 0, dupOrSliceIaeCtrHandle);
+        var throwDupOrSliceHandleWithBBArg = MethodHandles.dropArguments(throwDupOrSliceHandle, 0, ByteBuffer.class);
+
+        // construct a method handle that conditionally invokes 'bb.cleaner().clean()' or throws an
+        // IAE depending on the fact if the provided ByteBuffer has an attachment or not
+        var cleanIfNotAttachmentHandle = MethodHandles.guardWithTest(
+          isAttachmentNullHandle,
+          cleanBufferHandle,
+          throwDupOrSliceHandleWithBBArg);
         return buffer -> {
           try {
-            cleanBufferHandle.invokeExact(buffer);
+            cleanIfNotAttachmentHandle.invokeExact(buffer);
+          } catch (IllegalArgumentException exception) {
+            throw exception;
           } catch (Throwable _) {
           }
         };
       } catch (Throwable throwable) {
         UnsafeLogUtil.debug("Unable to access byte buffer cleaning methods; falling back to no cleaning", throwable);
-        return _ -> { // unable to clean direct buffers
-        };
+        return BB_CLEANER_NOOP; // unable to clean direct buffers
       }
     });
   }
@@ -1264,12 +1267,6 @@ public static void unsafeInvokeCleaner(ByteBuffer buffer) {
       throw new IllegalArgumentException("buffer is non-direct"); // mimics current behavior
     }
 
-    var attachmentGetter = BB_ATTACHMENT_GETTER.get();
-    var attachment = attachmentGetter.apply(buffer);
-    if (attachment != null) {
-      throw new IllegalArgumentException("duplicate or slice"); // mimics current behavior
-    }
-
     var cleanerInvoker = BB_CLEANER.get();
     cleanerInvoker.accept(buffer);
   }

wrapper-jvm/impl/src/test/java/eu/cloudnetservice/wrapper/impl/transform/unsafe/UnsafeReplacementDelegateTest.java

@@ -923,10 +923,30 @@ void testDefineClass() {
   @Test
   void testInvokeCleaner() {
     // can't really test if it worked, but can at least test if the method handle init works
-    Assertions.assertThrows(
-      IllegalArgumentException.class,
-      () -> UnsafeReplacementDelegate.unsafeInvokeCleaner(ByteBuffer.allocate(1)));
-    Assertions.assertDoesNotThrow(() -> UnsafeReplacementDelegate.unsafeInvokeCleaner(ByteBuffer.allocateDirect(1)));
+    {
+      var buffer = ByteBuffer.allocate(1);
+      var thrown = Assertions.assertThrows(
+        IllegalArgumentException.class,
+        () -> UnsafeReplacementDelegate.unsafeInvokeCleaner(buffer));
+      Assertions.assertNotNull(thrown.getMessage());
+      Assertions.assertEquals("buffer is non-direct", thrown.getMessage());
+    }
+
+    {
+      var buffer = ByteBuffer.allocateDirect(5);
+      var bufferSlice = buffer.slice(2, 2);
+      var thrown = Assertions.assertThrows(
+        IllegalArgumentException.class,
+        () -> UnsafeReplacementDelegate.unsafeInvokeCleaner(bufferSlice));
+      Assertions.assertNotNull(thrown.getMessage());
+      Assertions.assertEquals("duplicate or slice", thrown.getMessage());
+      Assertions.assertDoesNotThrow(() -> UnsafeReplacementDelegate.unsafeInvokeCleaner(buffer));
+    }
+
+    {
+      var cleanerConsumer = UnsafeReplacementDelegate.BB_CLEANER.get();
+      Assertions.assertNotSame(UnsafeReplacementDelegate.BB_CLEANER_NOOP, cleanerConsumer);
+    }
   }
 
   @Test