-
Notifications
You must be signed in to change notification settings - Fork 195
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Keep weak references to eager resources in session (#229)
- Loading branch information
1 parent
c4498eb
commit f6024dd
Showing
7 changed files
with
228 additions
and
17 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
70 changes: 70 additions & 0 deletions
70
...flow-core/tensorflow-core-api/src/main/java/org/tensorflow/internal/WeakPointerScope.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,70 @@ | ||
| package org.tensorflow.internal; | ||
|
|
||
| import java.util.Collections; | ||
| import java.util.Set; | ||
| import java.util.WeakHashMap; | ||
| 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> | ||
| * | ||
| * <p>It this {@code pointer} was already attached to this scope, this method has no effect.</p> | ||
| * | ||
| * @param pointer pointer to attach | ||
| * @throws IllegalStateException if that scope has already been closed | ||
| */ | ||
| public void attach(Pointer pointer) { | ||
| checkScope(); | ||
| if (pointers.add(pointer)) { | ||
| 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 | ||
| * @throws IllegalStateException if that scope has already been closed | ||
| */ | ||
| public void detach(Pointer pointer) { | ||
| checkScope(); | ||
| if (pointers.remove(pointer)) { | ||
| pointer.releaseReference(); | ||
| } | ||
| } | ||
|
|
||
| @Override | ||
| public synchronized void close() { | ||
| checkScope(); | ||
| pointers.forEach(Pointer::releaseReference); | ||
| pointers = null; | ||
| } | ||
|
|
||
| private Set<Pointer> pointers = Collections.newSetFromMap(new WeakHashMap<>()); | ||
|
|
||
| private void checkScope() { | ||
| if (pointers == null) { | ||
| throw new IllegalStateException("Pointer scope has been closed"); | ||
| } | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
114 changes: 114 additions & 0 deletions
114
...-core/tensorflow-core-api/src/test/java/org/tensorflow/internal/WeakPointerScopeTest.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,114 @@ | ||
| package org.tensorflow.internal; | ||
|
|
||
| import static org.junit.jupiter.api.Assertions.assertEquals; | ||
| import static org.junit.jupiter.api.Assertions.assertFalse; | ||
| import static org.junit.jupiter.api.Assertions.assertNull; | ||
| import static org.junit.jupiter.api.Assertions.assertThrows; | ||
| import static org.junit.jupiter.api.Assertions.assertTrue; | ||
|
|
||
| import org.bytedeco.javacpp.IntPointer; | ||
| import org.bytedeco.javacpp.Pointer; | ||
| import org.junit.jupiter.api.Test; | ||
| import org.tensorflow.EagerSession; | ||
|
|
||
| public class WeakPointerScopeTest { | ||
|
|
||
| @Test | ||
| public void resourcesAttachedAreFreedOnScopeClose() { | ||
| Pointer pointer = new IntPointer(10L); | ||
| assertEquals(0, pointer.referenceCount()); | ||
|
|
||
| try (WeakPointerScope scope = new WeakPointerScope()) { | ||
| scope.attach(pointer); | ||
| assertEquals(1, pointer.referenceCount()); | ||
| } | ||
| assertTrue(pointer.isNull()); | ||
| } | ||
|
|
||
| @Test | ||
| public void resourcesDetachedAreNotFreedOnScopeCloseWhenRetained() { | ||
| Pointer pointer = new IntPointer(10L); | ||
|
|
||
| try (WeakPointerScope scope = new WeakPointerScope()) { | ||
| scope.attach(pointer); | ||
| scope.detach(pointer.retainReference()); | ||
| } | ||
| assertFalse(pointer.isNull()); | ||
| assertEquals(1, pointer.referenceCount()); | ||
| pointer.deallocate(); | ||
| } | ||
|
|
||
| @Test | ||
| public void resourcesDetachedAreFreedWhenNotRetained() { | ||
| Pointer pointer = new IntPointer(10L); | ||
|
|
||
| try (WeakPointerScope scope = new WeakPointerScope()) { | ||
| scope.attach(pointer); | ||
|
|
||
| scope.detach(pointer); | ||
| assertTrue(pointer.isNull()); | ||
| } | ||
| } | ||
|
|
||
| @Test | ||
| public void attachingResourceMoreThanOnceHasNoEffect() { | ||
| Pointer pointer = new IntPointer(10L); | ||
|
|
||
| try (WeakPointerScope scope = new WeakPointerScope()) { | ||
| scope.attach(pointer); | ||
| scope.attach(pointer); | ||
| assertEquals(1, pointer.referenceCount()); | ||
|
|
||
| Pointer pointerCopy = new Pointer(pointer); | ||
| assertEquals(1, pointerCopy.referenceCount()); | ||
| scope.attach(pointerCopy); | ||
| assertEquals(1, pointerCopy.referenceCount()); | ||
| } | ||
| assertTrue(pointer.isNull()); | ||
| } | ||
|
|
||
| @Test | ||
| public void detachingUnattachedResourceHasNoEffect() { | ||
| Pointer pointer = new IntPointer(10L); | ||
| pointer.retainReference(); | ||
| assertEquals(1, pointer.referenceCount()); | ||
|
|
||
| try (WeakPointerScope scope = new WeakPointerScope()) { | ||
| scope.detach(pointer); | ||
| assertEquals(1, pointer.referenceCount()); | ||
| } | ||
| assertFalse(pointer.isNull()); | ||
| pointer.deallocate(); | ||
| } | ||
|
|
||
| @Test | ||
| public void operationOnClosedScopeFails() { | ||
| Pointer pointer = new IntPointer(10L); | ||
| WeakPointerScope scope = new WeakPointerScope(); | ||
| scope.close(); | ||
|
|
||
| assertThrows(IllegalStateException.class, () -> scope.attach(pointer)); | ||
| assertThrows(IllegalStateException.class, () -> scope.detach(pointer)); | ||
| assertThrows(IllegalStateException.class, () -> scope.close()); | ||
|
|
||
| pointer.deallocate(); | ||
| } | ||
|
|
||
| @Test | ||
| public void attachingResourceDoesNotPreventItToBeGarbageCollected() throws InterruptedException { | ||
| try (WeakPointerScope scope = new WeakPointerScope()) { | ||
| Pointer pointer = new IntPointer(10L); | ||
| scope.attach(pointer); | ||
| System.gc(); | ||
| Thread.sleep(50); | ||
|
|
||
| long before = Pointer.totalBytes(); | ||
| pointer = null; | ||
| System.gc(); | ||
| Thread.sleep(50); | ||
| long after = Pointer.totalBytes(); | ||
|
|
||
| assertEquals(4 * 10L, before - after); | ||
| } | ||
| } | ||
| } |