diff --git a/src/main/java/stormpot/internal/PreciseLeakDetector.java b/src/main/java/stormpot/internal/PreciseLeakDetector.java index af807228..adcf231e 100644 --- a/src/main/java/stormpot/internal/PreciseLeakDetector.java +++ b/src/main/java/stormpot/internal/PreciseLeakDetector.java @@ -20,6 +20,7 @@ import java.lang.ref.ReferenceQueue; import java.util.ArrayList; import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.LongAdder; public final class PreciseLeakDetector { @@ -30,11 +31,11 @@ public final class PreciseLeakDetector { public PreciseLeakDetector() { referenceQueue = new ReferenceQueue<>(); leakedObjectCount = new LongAdder(); - refs = new IdentityHashSet(); + refs = new CountedPhantomRefHashSet(); } public void register(BSlot slot) { - PhantomReference ref = new PhantomReference<>(slot.obj, referenceQueue); + CountedPhantomRef ref = new CountedPhantomRef<>(slot.obj, referenceQueue); slot.leakCheck = ref; synchronized (refs) { refs.add(ref); @@ -74,4 +75,22 @@ public long countLeakedObjects() { accumulateLeaks(); return leakedObjectCount.sum(); } + + private static final class CountedPhantomRef extends PhantomReference { + private static final AtomicInteger COUNTER = new AtomicInteger(); + + private final int id; // This hides in alignment padding on most JVMs, including Lilliput. + + CountedPhantomRef(T referent, ReferenceQueue q) { + super(referent, q); + id = COUNTER.getAndIncrement(); + } + } + + private static final class CountedPhantomRefHashSet extends IdentityHashSet { + @Override + protected int hashOf(Object obj) { + return ((CountedPhantomRef) obj).id; + } + } }