Keep weak references to eager resources in session

tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/EagerOperation.java

@@ -53,8 +53,6 @@ class EagerOperation extends AbstractOperation {
     this.name = name;
     this.opHandle = opNativeHandle;
     this.outputHandles = outputNativeHandles;
-    session.attach(opNativeHandle);
-    session.attach(outputNativeHandles);
     this.outputTensors = new AtomicReferenceArray<>(outputNativeHandles.length);
   }
 

tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/EagerOperationBuilder.java

@@ -65,12 +65,7 @@ final class EagerOperationBuilder implements OperationBuilder {
   @Override
   public EagerOperation build() {
     TFE_TensorHandle[] tensorHandles = execute(opHandle, session);
-    EagerOperation operation =
-        new EagerOperation(session, opHandle, tensorHandles, type, name);
-    // Release our reference to the native op handle now that we transferred its
-    // ownership to the EagerOperation
-    session.detach(opHandle);
-    return operation;
+    return new EagerOperation(session, opHandle, tensorHandles, type, name);
   }
 
   @Override

tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/EagerSession.java

@@ -24,6 +24,7 @@
 import org.bytedeco.javacpp.BytePointer;
 import org.bytedeco.javacpp.Pointer;
 import org.bytedeco.javacpp.PointerScope;
+import org.tensorflow.internal.WeakPointerScope;
 import org.tensorflow.internal.c_api.TFE_Context;
 import org.tensorflow.internal.c_api.TFE_ContextOptions;
 import org.tensorflow.internal.c_api.TF_Status;
@@ -326,14 +327,12 @@ void detach(Pointer... resources) {
 
   private static volatile EagerSession defaultSession = null;
 
-  private final PointerScope nativeResources;
+  private final WeakPointerScope nativeResources;
   private TFE_Context nativeHandle;
 
   private EagerSession(Options options) {
-    try (PointerScope scope = new PointerScope()) {
-      this.nativeResources = scope.extend();
-      this.nativeHandle = allocate(options.async, options.devicePlacementPolicy.code, options.config);
-    }
+    this.nativeResources = new WeakPointerScope();
+    this.nativeHandle = allocate(options.async, options.devicePlacementPolicy.code, options.config);
   }
 
   private void checkSession() {
@@ -363,7 +362,7 @@ private static TFE_Context allocate(boolean async, int devicePlacementPolicy, Co
       TFE_ContextOptionsSetDevicePlacementPolicy(opts, devicePlacementPolicy);
       TFE_Context context = TFE_NewContext(opts, status);
       status.throwExceptionIfNotOK();
-      return context;
+      return context.retainReference();
     }
   }
 

tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/internal/WeakPointerScope.java

@@ -0,0 +1,58 @@
+package org.tensorflow.internal;
+
+import java.lang.ref.WeakReference;
+import java.util.Collections;
+import java.util.LinkedList;
+import java.util.Set;
+import java.util.WeakHashMap;
+import java.util.concurrent.ConcurrentLinkedQueue;
+import org.bytedeco.javacpp.Pointer;
+
+/**
+ * A minimalist pointer scope only keeping weak references to its elements.
+ *
+ * <p>As opposed to {@link org.bytedeco.javacpp.PointerScope}, instances of this class will not
+ * prevent the garbage collector to free the memory of a pointer that is no longer reachable, even
+ * if it has been attached to the scope.</p>
+ *
+ * <p>When the scope is closed, all pointers that are still valid will be automatically deallocated
+ * while those already garbage-collected will be ignored.</p>
+ */
+public class WeakPointerScope implements AutoCloseable {
+
+  /**
+   * Attach a pointer to this scope.
+   *
+   * <p>Pointers attached to the scope will be automatically freed once the scope is closed, unless
+   * they have been already released by the garbage collector</p>
+   *
+   * @param pointer pointer to attach
+   */
+  public void attach(Pointer pointer) {
+    pointers.add(pointer.retainReference());
+  }
+
+  /**
+   * Detach a pointer from this scope.
+   *
+   * <p>Detaching a pointer from the scope will prevent its memory to be freed when closing the
+   * scope.</p>
+   *
+   * <p>If this {@code pointer} is not attached to this scope, this method has no effect.</p>
+   *
+   * @param pointer pointer to detach
+   */
+  public void detach(Pointer pointer) {
+    if (pointers.remove(pointer)) {
+      pointer.releaseReference();
+    }
+  }
+
+  @Override
+  public synchronized void close() {
+    pointers.forEach(Pointer::releaseReference);
+    pointers.clear();
+  }
+
+  private final Set<Pointer> pointers = Collections.newSetFromMap(new WeakHashMap<>());
+}

tensorflow-core/tensorflow-core-api/src/test/java/org/tensorflow/EagerOperationTest.java

@@ -35,8 +35,12 @@ public class EagerOperationTest {
   public void failToCreateIfSessionIsClosed() {
     EagerSession session = EagerSession.create();
     session.close();
-    try {
-      new EagerOperation(session, null, null, "Add", "add");
+    try (TInt32 t = TInt32.tensorOf(Shape.of(2, 3))) {
+      EagerOperation op =
+          opBuilder(session, "Const", "OutputAttrs")
+              .setAttr("dtype", t.dataType())
+              .setAttr("value", t)
+              .build();
       fail();
     } catch (IllegalStateException e) {
       // expected

tensorflow-core/tensorflow-core-api/src/test/java/org/tensorflow/EagerSessionTest.java

@@ -59,7 +59,6 @@ public void cleanupResourceInBackground() {
       sleep(50); // allow some time to the background thread for cleaning up resources
 
       long before = Pointer.totalBytes();
-      s.detach(ref.retainReference());
       ref = null;
       System.gc();
       sleep(50); // allow some time to the background thread for cleaning up resources