[GR-33950] Allocate and manage WasmMemory as a direct ByteBuffer.

PullRequest: graal/10132

Author: woess

wasm/src/org.graalvm.wasm/src/org/graalvm/wasm/WasmLanguage.java

@@ -44,7 +44,6 @@
 
 import org.graalvm.options.OptionDescriptors;
 import org.graalvm.wasm.api.WebAssembly;
-import org.graalvm.wasm.memory.UnsafeWasmMemory;
 import org.graalvm.wasm.memory.WasmMemory;
 
 import com.oracle.truffle.api.CallTarget;
@@ -105,9 +104,7 @@ protected void finalizeContext(WasmContext context) {
         super.finalizeContext(context);
         for (int i = 0; i < context.memories().count(); ++i) {
             final WasmMemory memory = context.memories().memory(i);
-            if (memory instanceof UnsafeWasmMemory) {
-                ((UnsafeWasmMemory) memory).free();
-            }
+            memory.close();
         }
         try {
             context.fdManager().close();

wasm/src/org.graalvm.wasm/src/org/graalvm/wasm/api/WebAssembly.java

@@ -40,13 +40,17 @@
  */
 package org.graalvm.wasm.api;
 
-import com.oracle.truffle.api.CompilerAsserts;
-import com.oracle.truffle.api.TruffleContext;
-import com.oracle.truffle.api.interop.InteropException;
-import com.oracle.truffle.api.interop.InteropLibrary;
-import com.oracle.truffle.api.interop.InvalidArrayIndexException;
-import com.oracle.truffle.api.interop.UnknownIdentifierException;
-import com.oracle.truffle.api.interop.UnsupportedMessageException;
+import static java.lang.Integer.compareUnsigned;
+import static org.graalvm.wasm.WasmMath.minUnsigned;
+import static org.graalvm.wasm.api.JsConstants.JS_LIMITS;
+
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+
 import org.graalvm.collections.EconomicMap;
 import org.graalvm.collections.Pair;
 import org.graalvm.wasm.EmbedderDataHolder;
@@ -73,15 +77,13 @@
 import org.graalvm.wasm.memory.UnsafeWasmMemory;
 import org.graalvm.wasm.memory.WasmMemory;
 
-import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-import java.util.Objects;
-
-import static java.lang.Integer.compareUnsigned;
-import static org.graalvm.wasm.WasmMath.minUnsigned;
-import static org.graalvm.wasm.api.JsConstants.JS_LIMITS;
+import com.oracle.truffle.api.CompilerAsserts;
+import com.oracle.truffle.api.TruffleContext;
+import com.oracle.truffle.api.interop.InteropException;
+import com.oracle.truffle.api.interop.InteropLibrary;
+import com.oracle.truffle.api.interop.InvalidArrayIndexException;
+import com.oracle.truffle.api.interop.UnknownIdentifierException;
+import com.oracle.truffle.api.interop.UnsupportedMessageException;
 
 public class WebAssembly extends Dictionary {
     private final WasmContext currentContext;
@@ -103,6 +105,7 @@ public WebAssembly(WasmContext currentContext) {
         addMember("mem_alloc", new Executable(args -> memAlloc(args)));
         addMember("mem_grow", new Executable(args -> memGrow(args)));
         addMember("mem_set_grow_callback", new Executable(args -> memSetGrowCallback(args)));
+        addMember("mem_as_byte_buffer", new Executable(args -> memAsByteBuffer(args)));
 
         addMember("global_alloc", new Executable(args -> globalAlloc(args)));
         addMember("global_read", new Executable(args -> globalRead(args)));
@@ -620,6 +623,18 @@ public static void invokeMemGrowCallback(WasmMemory memory) {
         }
     }
 
+    private static Object memAsByteBuffer(Object[] args) {
+        checkArgumentCount(args, 1);
+        if (args[0] instanceof WasmMemory) {
+            WasmMemory memory = (WasmMemory) args[0];
+            ByteBuffer buffer = memory.asByteBuffer();
+            if (buffer != null) {
+                return WasmContext.get(null).environment().asGuestValue(buffer);
+            }
+        }
+        return WasmVoidResult.getInstance();
+    }
+
     private static Object globalAlloc(Object[] args) {
         checkArgumentCount(args, 3);
 

wasm/src/org.graalvm.wasm/src/org/graalvm/wasm/memory/ByteArrayWasmMemory.java

@@ -47,6 +47,8 @@
 import static org.graalvm.wasm.constants.Sizes.MAX_MEMORY_INSTANCE_SIZE;
 import static org.graalvm.wasm.constants.Sizes.MEMORY_PAGE_SIZE;
 
+import java.nio.ByteBuffer;
+
 import org.graalvm.wasm.constants.Sizes;
 import org.graalvm.wasm.exception.Failure;
 import org.graalvm.wasm.exception.WasmException;
@@ -408,4 +410,14 @@ public WasmMemory duplicate() {
         System.arraycopy(buffer, 0, other.buffer, 0, buffer.length);
         return other;
     }
+
+    @Override
+    public void close() {
+        buffer = null;
+    }
+
+    @Override
+    public ByteBuffer asByteBuffer() {
+        return null;
+    }
 }

wasm/src/org.graalvm.wasm/src/org/graalvm/wasm/memory/UnsafeWasmMemory.java

@@ -48,22 +48,19 @@
 import static org.graalvm.wasm.constants.Sizes.MEMORY_PAGE_SIZE;
 
 import java.lang.reflect.Field;
+import java.nio.ByteBuffer;
 
 import org.graalvm.wasm.constants.Sizes;
 import org.graalvm.wasm.exception.Failure;
 import org.graalvm.wasm.exception.WasmException;
 
 import com.oracle.truffle.api.CompilerDirectives;
 import com.oracle.truffle.api.CompilerDirectives.TruffleBoundary;
-import com.oracle.truffle.api.interop.InteropLibrary;
-import com.oracle.truffle.api.library.ExportLibrary;
-import com.oracle.truffle.api.library.ExportMessage;
 import com.oracle.truffle.api.nodes.Node;
 
 import sun.misc.Unsafe;
 
-@ExportLibrary(InteropLibrary.class)
-public final class UnsafeWasmMemory extends WasmMemory implements AutoCloseable {
+public final class UnsafeWasmMemory extends WasmMemory {
     /**
      * @see #declaredMinSize()
      */
@@ -88,6 +85,8 @@ public final class UnsafeWasmMemory extends WasmMemory implements AutoCloseable
      */
     private final int maxAllowedSize;
 
+    private ByteBuffer buffer;
+
     private static final Unsafe unsafe;
 
     private UnsafeWasmMemory(int declaredMinSize, int declaredMaxSize, int initialSize, int maxAllowedSize) {
@@ -102,13 +101,23 @@ private UnsafeWasmMemory(int declaredMinSize, int declaredMaxSize, int initialSi
         this.size = declaredMinSize;
         this.maxAllowedSize = maxAllowedSize;
         final long byteSize = byteSize();
+        this.buffer = allocateBuffer(byteSize);
+        this.startAddress = getBufferAddress(buffer);
+    }
+
+    @TruffleBoundary
+    private static ByteBuffer allocateBuffer(final long byteSize) {
+        assert (int) byteSize == byteSize : byteSize;
         try {
-            this.startAddress = unsafe.allocateMemory(byteSize);
+            return ByteBuffer.allocateDirect((int) byteSize);
         } catch (OutOfMemoryError error) {
-            CompilerDirectives.transferToInterpreter();
             throw WasmException.create(Failure.MEMORY_ALLOCATION_FAILED);
         }
-        unsafe.setMemory(startAddress, byteSize, (byte) 0);
+    }
+
+    @TruffleBoundary
+    private static long getBufferAddress(ByteBuffer buffer) {
+        return ((sun.nio.ch.DirectBuffer) buffer).address();
     }
 
     public UnsafeWasmMemory(int declaredMinSize, int declaredMaxSize, int maxAllowedSize) {
@@ -170,10 +179,10 @@ public boolean grow(int extraPageSize) {
             // computation of targetByteSize does not overflow.
             final int targetByteSize = multiplyExact(addExact(size(), extraPageSize), MEMORY_PAGE_SIZE);
             try {
-                final long updatedStartAddress = unsafe.allocateMemory(targetByteSize);
+                ByteBuffer updatedBuffer = allocateBuffer(targetByteSize);
+                final long updatedStartAddress = getBufferAddress(updatedBuffer);
                 unsafe.copyMemory(startAddress, updatedStartAddress, byteSize());
-                unsafe.setMemory(updatedStartAddress + byteSize(), targetByteSize - byteSize(), (byte) 0);
-                unsafe.freeMemory(startAddress);
+                buffer = updatedBuffer;
                 startAddress = updatedStartAddress;
                 size += extraPageSize;
                 invokeGrowCallback();
@@ -348,7 +357,7 @@ public WasmMemory duplicate() {
     }
 
     public void free() {
-        unsafe.freeMemory(this.startAddress);
+        buffer = null;
         startAddress = 0;
         size = 0;
     }
@@ -364,15 +373,9 @@ public void close() {
         }
     }
 
-    @SuppressWarnings("static-method")
-    @ExportMessage
-    public boolean isPointer() {
-        return true;
-    }
-
-    @ExportMessage
-    public long asPointer() {
-        return startAddress;
+    @Override
+    public ByteBuffer asByteBuffer() {
+        return buffer.duplicate();
     }
 
     static {

wasm/src/org.graalvm.wasm/src/org/graalvm/wasm/memory/WasmMemory.java

@@ -40,6 +40,7 @@
  */
 package org.graalvm.wasm.memory;
 
+import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
 import java.nio.charset.StandardCharsets;
 
@@ -68,7 +69,7 @@
 import com.oracle.truffle.api.profiles.BranchProfile;
 
 @ExportLibrary(InteropLibrary.class)
-public abstract class WasmMemory extends EmbedderDataHolder implements TruffleObject {
+public abstract class WasmMemory extends EmbedderDataHolder implements TruffleObject, AutoCloseable {
 
     private Object growCallback = null;
 
@@ -505,4 +506,9 @@ public Object getGrowCallback() {
     protected void invokeGrowCallback() {
         WebAssembly.invokeMemGrowCallback(this);
     }
+
+    @Override
+    public abstract void close();
+
+    public abstract ByteBuffer asByteBuffer();
 }