SwiftArena.ofAuto() which uses GC to manage swift instances (and destroy them) (swiftlang#133)

Author: ktoso

SwiftKit/src/main/java/org/swift/swiftkit/AutoSwiftMemorySession.java

@@ -0,0 +1,70 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2024 Apple Inc. and the Swift.org project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of Swift.org project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+package org.swift.swiftkit;
+
+import java.lang.foreign.MemorySegment;
+import java.lang.ref.Cleaner;
+import java.util.Objects;
+import java.util.concurrent.ThreadFactory;
+
+/**
+ * A memory session which manages registered objects via the Garbage Collector.
+ *
+ * <p> When registered Java wrapper classes around native Swift instances {@link SwiftInstance},
+ * are eligible for collection, this will trigger the cleanup of the native resources as well.
+ *
+ * <p> This memory session is LESS reliable than using a {@link ConfinedSwiftMemorySession} because
+ * the timing of when the native resources are cleaned up is somewhat undefined, and rely on the
+ * system GC. Meaning, that if an object nas been promoted to an old generation, there may be a
+ * long time between the resource no longer being referenced "in Java" and its native memory being released,
+ * and also the deinit of the Swift type being run.
+ *
+ * <p> This can be problematic for Swift applications which rely on quick release of resources, and may expect
+ * the deinits to run in expected and "quick" succession.
+ *
+ * <p> Whenever possible, prefer using an explicitly managed {@link SwiftArena}, such as {@link SwiftArena#ofConfined()}.
+ */
+final class AutoSwiftMemorySession implements SwiftArena {
+
+    private final Cleaner cleaner;
+
+    public AutoSwiftMemorySession(ThreadFactory cleanerThreadFactory) {
+        this.cleaner = Cleaner.create(cleanerThreadFactory);
+    }
+
+    @Override
+    public void register(SwiftHeapObject object) {
+        SwiftHeapObjectCleanup cleanupAction = new SwiftHeapObjectCleanup(object.$memorySegment(), object.$swiftType());
+        register(object, cleanupAction);
+    }
+
+    // visible for testing
+    void register(SwiftHeapObject object, SwiftHeapObjectCleanup cleanupAction) {
+        Objects.requireNonNull(object, "obj");
+        Objects.requireNonNull(cleanupAction, "cleanupAction");
+
+
+        cleaner.register(object, cleanupAction);
+    }
+
+    @Override
+    public void register(SwiftValue value) {
+        Objects.requireNonNull(value, "value");
+        MemorySegment resource = value.$memorySegment();
+        var cleanupAction = new SwiftValueCleanup(resource);
+        cleaner.register(value, cleanupAction);
+    }
+
+}

SwiftKit/src/main/java/org/swift/swiftkit/ClosableSwiftArena.java

@@ -0,0 +1,28 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2024 Apple Inc. and the Swift.org project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of Swift.org project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+package org.swift.swiftkit;
+
+/**
+ * Auto-closable version of {@link SwiftArena}.
+ */
+public interface ClosableSwiftArena extends SwiftArena, AutoCloseable {
+
+    /**
+     * Close the arena and make sure all objects it managed are released.
+     * Throws if unable to verify all resources have been release (e.g. over retained Swift classes)
+     */
+    void close();
+
+}

SwiftKit/src/main/java/org/swift/swiftkit/ConfinedSwiftMemorySession.java

@@ -0,0 +1,86 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2024 Apple Inc. and the Swift.org project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of Swift.org project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+package org.swift.swiftkit;
+
+import java.util.LinkedList;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicInteger;
+
+final class ConfinedSwiftMemorySession implements ClosableSwiftArena {
+
+    final static int CLOSED = 0;
+    final static int ACTIVE = 1;
+
+    final Thread owner;
+    final AtomicInteger state;
+
+    final ConfinedResourceList resources;
+
+    public ConfinedSwiftMemorySession(Thread owner) {
+        this.owner = owner;
+        this.state = new AtomicInteger(ACTIVE);
+        this.resources = new ConfinedResourceList();
+    }
+
+    public void checkValid() throws RuntimeException {
+        if (this.owner != null && this.owner != Thread.currentThread()) {
+            throw new WrongThreadException("ConfinedSwift arena is confined to %s but was closed from %s!".formatted(this.owner, Thread.currentThread()));
+        } else if (this.state.get() < ACTIVE) {
+            throw new RuntimeException("SwiftArena is already closed!");
+        }
+    }
+
+    @Override
+    public void close() {
+        checkValid();
+
+        // Cleanup all resources
+        if (this.state.compareAndExchange(ACTIVE, CLOSED) == ACTIVE) {
+            this.resources.runCleanup();
+        } // else, was already closed; do nothing
+    }
+
+    @Override
+    public void register(SwiftHeapObject object) {
+        checkValid();
+
+        var cleanup = new SwiftHeapObjectCleanup(object.$memorySegment(), object.$swiftType());
+        this.resources.add(cleanup);
+    }
+
+    @Override
+    public void register(SwiftValue value) {
+        checkValid();
+
+        var cleanup = new SwiftValueCleanup(value.$memorySegment());
+        this.resources.add(cleanup);
+    }
+
+    static final class ConfinedResourceList implements SwiftResourceList {
+        // TODO: Could use intrusive linked list to avoid one indirection here
+        final List<SwiftInstanceCleanup> resourceCleanups = new LinkedList<>();
+
+        void add(SwiftInstanceCleanup cleanup) {
+            resourceCleanups.add(cleanup);
+        }
+
+        @Override
+        public void runCleanup() {
+            for (SwiftInstanceCleanup cleanup : resourceCleanups) {
+                cleanup.run();
+            }
+        }
+    }
+}

SwiftKit/src/main/java/org/swift/swiftkit/SwiftArena.java

@@ -14,23 +14,26 @@
 
 package org.swift.swiftkit;
 
-import java.util.LinkedList;
-import java.util.List;
-import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.ThreadFactory;
 
 /**
  * A Swift arena manages Swift allocated memory for classes, structs, enums etc.
  * When an arena is closed, it will destroy all managed swift objects in a way appropriate to their type.
- * <p>
- * A confined arena has an associated owner thread that confines some operations to
- * associated owner thread such as {@link #close()}.
+ *
+ * <p> A confined arena has an associated owner thread that confines some operations to
+ * associated owner thread such as {@link ClosableSwiftArena#close()}.
  */
-public interface SwiftArena extends AutoCloseable {
+public interface SwiftArena {
 
-    static SwiftArena ofConfined() {
+    static ClosableSwiftArena ofConfined() {
         return new ConfinedSwiftMemorySession(Thread.currentThread());
     }
 
+    static SwiftArena ofAuto() {
+        ThreadFactory cleanerThreadFactory = r -> new Thread(r, "AutoSwiftArenaCleanerThread");
+        return new AutoSwiftMemorySession(cleanerThreadFactory);
+    }
+
     /**
      * Register a Swift reference counted heap object with this arena (such as a {@code class} or {@code actor}).
      * Its memory should be considered managed by this arena, and be destroyed when the arena is closed.
@@ -43,95 +46,23 @@ static SwiftArena ofConfined() {
      */
     void register(SwiftValue value);
 
-    /**
-     * Close the arena and make sure all objects it managed are released.
-     * Throws if unable to verify all resources have been release (e.g. over retained Swift classes)
-     */
-    void close();
-
 }
 
-final class ConfinedSwiftMemorySession implements SwiftArena {
-
-//    final Arena underlying;
-    final Thread owner;
-    final SwiftResourceList resources;
-
-    final int CLOSED = 0;
-    final int ACTIVE = 1;
-    final AtomicInteger state;
-
-    public ConfinedSwiftMemorySession(Thread owner) {
-        this.owner = owner;
-        resources = new ConfinedResourceList();
-        state = new AtomicInteger(ACTIVE);
-    }
-
-    public void checkValid() throws RuntimeException {
-        if (this.owner != null && this.owner != Thread.currentThread()) {
-            throw new WrongThreadException("ConfinedSwift arena is confined to %s but was closed from %s!".formatted(this.owner, Thread.currentThread()));
-        } else if (this.state.get() < ACTIVE) {
-            throw new RuntimeException("Arena is already closed!");
-        }
-    }
-
-    @Override
-    public void register(SwiftHeapObject object) {
-        this.resources.add(new SwiftHeapObjectCleanup(object));
-    }
-
-    @Override
-    public void register(SwiftValue value) {
-        this.resources.add(new SwiftValueCleanup(value.$memorySegment()));
-    }
-
-    @Override
-    public void close() {
-        checkValid();
-
-        // Cleanup all resources
-        if (this.state.compareAndExchange(ACTIVE, CLOSED) == ACTIVE) {
-            this.resources.cleanup();
-        } // else, was already closed; do nothing
-    }
-
-    /**
-     * Represents a list of resources that need a cleanup, e.g. allocated classes/structs.
-     */
-    static abstract class SwiftResourceList implements Runnable {
-        // TODO: Could use intrusive linked list to avoid one indirection here
-        final List<SwiftInstanceCleanup> resourceCleanups = new LinkedList<>();
-
-        abstract void add(SwiftInstanceCleanup cleanup);
-
-        public abstract void cleanup();
-
-        public final void run() {
-            cleanup(); // cleaner interop
-        }
-    }
-
-    static final class ConfinedResourceList extends SwiftResourceList {
-        @Override
-        void add(SwiftInstanceCleanup cleanup) {
-            resourceCleanups.add(cleanup);
-        }
-
-        @Override
-        public void cleanup() {
-            for (SwiftInstanceCleanup cleanup : resourceCleanups) {
-                cleanup.run();
-            }
-        }
-    }
+/**
+ * Represents a list of resources that need a cleanup, e.g. allocated classes/structs.
+ */
+interface SwiftResourceList {
 
+    void runCleanup();
 
 }
 
+
 final class UnexpectedRetainCountException extends RuntimeException {
     public UnexpectedRetainCountException(Object resource, long retainCount, int expectedRetainCount) {
         super(("Attempting to cleanup managed memory segment %s, but it's retain count was different than [%d] (was %d)! " +
                 "This would result in destroying a swift object that is still retained by other code somewhere."
         ).formatted(resource, expectedRetainCount, retainCount));
     }
 }
+

SwiftKit/src/main/java/org/swift/swiftkit/SwiftInstanceCleanup.java

@@ -19,23 +19,43 @@
 /**
  * A Swift memory instance cleanup, e.g. count-down a reference count and destroy a class, or destroy struct/enum etc.
  */
-sealed interface SwiftInstanceCleanup extends Runnable {
+interface SwiftInstanceCleanup extends Runnable {
 }
 
-record SwiftHeapObjectCleanup(SwiftHeapObject instance) implements SwiftInstanceCleanup {
+/**
+ * Implements cleaning up a Swift {@link SwiftHeapObject}.
+ * <p>
+ * This class does not store references to the Java wrapper class, and therefore the wrapper may be subject to GC,
+ * which may trigger a cleanup (using this class), which will clean up its underlying native memory resource.
+ */
+// non-final for testing
+class SwiftHeapObjectCleanup implements SwiftInstanceCleanup {
+
+    final MemorySegment selfPointer;
+    final SwiftAnyType selfType;
+
+    /**
+     * This constructor on purpose does not just take a {@link SwiftHeapObject} in order to make it very
+     * clear that it does not take ownership of it, but we ONLY manage the native resource here.
+     *
+     * This is important for {@link AutoSwiftMemorySession} which relies on the wrapper type to be GC-able,
+     * when no longer "in use" on the Java side.
+     */
+    SwiftHeapObjectCleanup(MemorySegment selfPointer, SwiftAnyType selfType) {
+        this.selfPointer  = selfPointer;
+        this.selfType = selfType;
+    }
 
     @Override
     public void run() throws UnexpectedRetainCountException {
         // Verify we're only destroying an object that's indeed not retained by anyone else:
-        long retainedCount = SwiftKit.retainCount(this.instance);
+        long retainedCount = SwiftKit.retainCount(selfPointer);
         if (retainedCount > 1) {
-            throw new UnexpectedRetainCountException(this.instance, retainedCount, 1);
+            throw new UnexpectedRetainCountException(selfPointer, retainedCount, 1);
         }
 
         // Destroy (and deinit) the object:
-        var ty = this.instance.$swiftType();
-
-        SwiftValueWitnessTable.destroy(ty, this.instance.$memorySegment());
+        SwiftValueWitnessTable.destroy(selfType, selfPointer);
 
         // Invalidate the Java wrapper class, in order to prevent effectively use-after-free issues.
         // FIXME: some trouble with setting the pointer to null, need to figure out an appropriate way to do this