memory: Improve SharedMemory handling, cleaner

Author: kohlschuetter

junixsocket-memory/src/main/java/org/newsclub/net/unix/memory/SharedMemory.java

@@ -102,7 +102,7 @@ private SharedMemory(FileDescriptor fd, long size, String name, int mopts) throw
       this.size = size;
     }
     Objects.requireNonNull(fd);
-    this.cleaner = new SharedMemoryCleaner(this, fd);
+    this.cleaner = new SharedMemoryCleaner(null, this, fd);
     this.name = name;
     this.knownReadOnly = isReadOnly(mopts);
     this.unlinkUponClose = (mopts & MemoryImplUtilInternal.MOPT_UNLINK_UPON_CLOSE) != 0;
@@ -459,11 +459,12 @@ public MemorySegment asMappedMemorySegment(MapMode mapMode, Arena arena, long of
       throw new IllegalArgumentException("length");
     }
 
+    // use a zero-length, 0-address segment for lifecycle management, preventing chicken-egg problem
+    MemorySegment arenaSegment;
     if (arena == null) {
-      arena = cleaner.defaultArena;
-      if (arena == null) {
-        arena = cleaner.defaultArena = Arena.ofShared();
-      }
+      arenaSegment = cleaner.getArenaSegment();
+    } else {
+      arenaSegment = arena.allocate(0);
     }
 
     int mmode = resolveMmode(mapMode);
@@ -474,11 +475,19 @@ public MemorySegment asMappedMemorySegment(MapMode mapMode, Arena arena, long of
       length = size;
     }
 
-    // use a zero-length, 0-address segment for lifecycle management, preventing chicken-egg problem
-    MemorySegment arenaSegment = arena.allocate(0);
+    ByteBuffer buf = getUtil().mmapShm(arenaSegment, cleaner.fd, offset, length, mmode, duplicates);
+    return asRegisteredMemorySegment(cleaner, buf, (mmode
+        & MemoryImplUtilInternal.MMODE_WRITE) != 0, duplicates);
+  }
+
+  static MemorySegment asRegisteredMemorySegment(SharedMemoryCleaner cleaner, ByteBuffer buf,
+      boolean rw) {
+    return asRegisteredMemorySegment(cleaner, buf, rw, 0);
+  }
 
-    ByteBuffer buf = getUtil().mmap(arenaSegment, cleaner.fd, offset, length, mmode, duplicates);
-    if ((mmode & MemoryImplUtilInternal.MMODE_WRITE) == 0) {
+  private static MemorySegment asRegisteredMemorySegment(SharedMemoryCleaner cleaner,
+      ByteBuffer buf, boolean rw, int duplicates) {
+    if (!rw) {
       // The MapMode is read-only.
       // If we don't ask for a read-only buffer here, write accesses will fail with a page fault
       // ("java.lang.InternalError: a fault occurred in an unsafe memory access operation")

junixsocket-memory/src/main/java/org/newsclub/net/unix/memory/SharedMemoryCleaner.java

@@ -21,8 +21,11 @@
 import java.io.IOException;
 import java.lang.foreign.Arena;
 import java.lang.foreign.MemorySegment;
+import java.util.ArrayList;
 import java.util.LinkedHashMap;
+import java.util.List;
 import java.util.Map;
+import java.util.Map.Entry;
 import java.util.Objects;
 import java.util.WeakHashMap;
 
@@ -32,13 +35,24 @@
 class SharedMemoryCleaner extends CleanableState {
   private final Map<MemorySegment, Integer> segments = new LinkedHashMap<>();
   private final Map<Futex, Futex> futexes = new WeakHashMap<>();
-  private final Arena arena = Arena.ofShared();
-  private final MemorySegment arenaSegment = arena.allocate(0);
+  private Arena arena;
+  private final boolean closeArena;
+  private MemorySegment arenaSegment;
   final FileDescriptor fd;
-  Arena defaultArena;
 
-  protected SharedMemoryCleaner(Object observed, FileDescriptor fd) {
+  SharedMemoryCleaner(Arena arena, Object observed) {
+    this(arena, observed, null);
+  }
+
+  SharedMemoryCleaner(Arena arena, Object observed, FileDescriptor fd) {
     super(observed);
+    if (arena == null) {
+      this.arena = null;
+      this.closeArena = true;
+    } else {
+      this.arena = arena;
+      this.closeArena = false;
+    }
     this.fd = fd;
   }
 
@@ -50,14 +64,14 @@ void registerMemorySegment(MemorySegment ms, int duplicates) {
 
   @Override
   @SuppressWarnings("PMD.CognitiveComplexity")
-  protected void doClean() throws IOException {
+  protected synchronized void doClean() throws IOException {
     if (!SharedMemory.isUtilLoaded()) {
       // Nothing to do
       return;
     }
 
     Map<FileDescriptor, Long> map = SharedMemory.FD_MEMORY;
-    if (map != null) {
+    if (map != null && fd != null) {
       synchronized (map) {
         map.remove(fd);
       }
@@ -77,45 +91,46 @@ protected void doClean() throws IOException {
       }
     }
 
-    if (defaultArena != null) {
-      defaultArena.close();
-      defaultArena = null;
+    if (closeArena && arena != null) {
+      arena.close();
     }
 
     MemoryImplUtilInternal util = SharedMemory.getUtil();
 
     IOException exc = null;
-    if (fd.valid()) {
+    if (fd != null && fd.valid()) {
       try {
         util.close(fd);
       } catch (IOException e) {
         exc = e;
       }
     }
 
+    List<Entry<MemorySegment, Integer>> list;
     synchronized (segments) {
-      for (Map.Entry<MemorySegment, Integer> en : segments.entrySet()) {
-        MemorySegment ms = en.getKey();
-        int duplicates = en.getValue();
-
-        long addr = ms.address();
-        long length = ms.byteSize();
-        if (ms.scope().isAlive()) {
-          util.madvise(addr, length, MemoryImplUtilInternal.MADV_FREE_NOW, true);
-          continue;
-        }
+      list = new ArrayList<>(segments.entrySet()).reversed();
+      segments.clear();
+    }
+    for (Map.Entry<MemorySegment, Integer> en : list) {
+      MemorySegment ms = en.getKey();
+      int duplicates = en.getValue();
+
+      long addr = ms.address();
+      long length = ms.byteSize();
+      if (ms.scope().isAlive()) {
+        util.madvise(addr, length, MemoryImplUtilInternal.MADV_FREE_NOW, true);
+        continue;
+      }
 
-        try {
-          util.unmap(addr, length, duplicates, false);
-        } catch (IOException e) {
-          if (exc == null) {
-            exc = e;
-          } else {
-            exc.addSuppressed(e);
-          }
+      try {
+        util.unmap(addr, length, duplicates, false);
+      } catch (IOException e) {
+        if (exc == null) {
+          exc = e;
+        } else {
+          exc.addSuppressed(e);
         }
       }
-      segments.clear();
     }
 
     if (exc != null) {
@@ -141,7 +156,10 @@ public boolean isCovered(MemorySegment segment) {
     return false;
   }
 
-  public MemorySegment getArenaSegment() {
+  public synchronized MemorySegment getArenaSegment() {
+    if (this.arenaSegment == null) {
+      this.arenaSegment = getArena().allocate(0);
+    }
     return arenaSegment;
   }
 
@@ -156,4 +174,11 @@ public void checkCovered(MemorySegment addr) throws IOException {
       throw new IOException("Not a MemorySegment of ours");
     }
   }
+
+  public synchronized Arena getArena() {
+    if (this.arena == null) {
+      this.arena = Arena.ofShared();
+    }
+    return this.arena;
+  }
 }