Skip to content

Commit

Permalink
Refactored mock cache to be non-blocking
Browse files Browse the repository at this point in the history
  • Loading branch information
@raphw
raphw committed Jun 27, 2016
1 parent ff033dc commit 0ff47d4
Show file tree
Hide file tree
Showing 3 changed files with 133 additions and 44 deletions.
2 changes: 1 addition & 1 deletion src/main/java/org/mockito/internal/creation/bytebuddy/ByteBuddyMockMaker.java
View file
Expand Up @@ -19,7 +19,7 @@ public class ByteBuddyMockMaker implements MockMaker {
private final CachingMockBytecodeGenerator cachingMockBytecodeGenerator;

public ByteBuddyMockMaker() {
cachingMockBytecodeGenerator = new CachingMockBytecodeGenerator();
cachingMockBytecodeGenerator = new CachingMockBytecodeGenerator(false);
}

@Override
Expand Down
170 changes: 128 additions & 42 deletions src/main/java/org/mockito/internal/creation/bytebuddy/CachingMockBytecodeGenerator.java
View file
@@ -1,70 +1,95 @@
package org.mockito.internal.creation.bytebuddy;

import static org.mockito.internal.util.StringJoiner.join;
import org.mockito.exceptions.base.MockitoException;

import java.lang.ref.Reference;
import java.lang.ref.ReferenceQueue;
import java.lang.ref.SoftReference;
import java.lang.ref.WeakReference;
import java.lang.reflect.Modifier;
import java.net.URL;
import java.net.URLClassLoader;
import java.util.Collections;
import java.util.HashSet;
import java.util.Set;
import java.util.WeakHashMap;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;
import org.mockito.exceptions.base.MockitoException;
import java.util.concurrent.ConcurrentMap;

import static org.mockito.internal.util.StringJoiner.join;

class CachingMockBytecodeGenerator {
class CachingMockBytecodeGenerator extends ReferenceQueue<ClassLoader> {

private final Lock avoidingClassLeakCacheLock = new ReentrantLock();
public final WeakHashMap<ClassLoader, CachedBytecodeGenerator> avoidingClassLeakageCache =
new WeakHashMap<ClassLoader, CachedBytecodeGenerator>();
private static final ClassLoader BOOT_LOADER = new URLClassLoader(new URL[0], null);

final ConcurrentMap<Key<ClassLoader>, CachedBytecodeGenerator> avoidingClassLeakageCache =
new ConcurrentHashMap<Key<ClassLoader>, CachedBytecodeGenerator>();

private final MockBytecodeGenerator mockBytecodeGenerator = new MockBytecodeGenerator();

private final boolean weak;

public CachingMockBytecodeGenerator(boolean weak) {
this.weak = weak;
}

@SuppressWarnings("unchecked")
public <T> Class<T> get(MockFeatures<T> params) {
// TODO improves locking behavior with ReentrantReadWriteLock ?
avoidingClassLeakCacheLock.lock();
try {

Class<?> generatedMockClass = mockCachePerClassLoaderOf(params.mockedType).getOrGenerateMockClass(
params
);
pollAndClearReferences();
return (Class<T>) mockCachePerClassLoaderOf(params.mockedType.getClassLoader()).getOrGenerateMockClass(params);
}

return (Class<T>) generatedMockClass;
} finally {
avoidingClassLeakCacheLock.unlock();
void pollAndClearReferences() {
Reference<?> reference;
while ((reference = poll()) != null) {
avoidingClassLeakageCache.remove(reference);
}
}

private <T> CachedBytecodeGenerator mockCachePerClassLoaderOf(Class<T> mockedType) {
if (!avoidingClassLeakageCache.containsKey(mockedType.getClassLoader())) {
avoidingClassLeakageCache.put(
mockedType.getClassLoader(),
new CachedBytecodeGenerator(mockBytecodeGenerator)
);
private CachedBytecodeGenerator mockCachePerClassLoaderOf(ClassLoader classLoader) {
classLoader = classLoader == null ? BOOT_LOADER : classLoader;
CachedBytecodeGenerator generator = avoidingClassLeakageCache.get(new StrongKey<ClassLoader>(classLoader));
if (generator == null) {
CachedBytecodeGenerator newGenerator = new CachedBytecodeGenerator(mockBytecodeGenerator, weak);
generator = avoidingClassLeakageCache.putIfAbsent(new WeakKey<ClassLoader>(classLoader, this), newGenerator);
if (generator == null) {
generator = newGenerator;
}
}
return avoidingClassLeakageCache.get(mockedType.getClassLoader());
return generator;
}

private static class CachedBytecodeGenerator {
private ConcurrentHashMap<MockKey, WeakReference<Class<?>>> generatedClassCache =
new ConcurrentHashMap<MockKey, WeakReference<Class<?>>>();

private ConcurrentHashMap<MockKey, Reference<Class<?>>> generatedClassCache = new ConcurrentHashMap<MockKey, Reference<Class<?>>>();
private final MockBytecodeGenerator generator;
private final boolean weak;

private CachedBytecodeGenerator(MockBytecodeGenerator generator) {
private CachedBytecodeGenerator(MockBytecodeGenerator generator, boolean weak) {
this.generator = generator;
this.weak = weak;
}

public <T> Class<?> getOrGenerateMockClass(MockFeatures<T> features) {
MockKey<?> mockKey = MockKey.of(features.mockedType, features.interfaces);
Class<?> generatedMockClass = null;
WeakReference<Class<?>> classWeakReference = generatedClassCache.get(mockKey);
if(classWeakReference != null) {
generatedMockClass = classWeakReference.get();
private Class<?> getMockClass(MockKey<?> mockKey) {
Reference<Class<?>> classReference = generatedClassCache.get(mockKey);
if(classReference != null) {
return classReference.get();
} else {
return null;
}
if(generatedMockClass == null) {
generatedMockClass = generate(features);
}

Class<?> getOrGenerateMockClass(MockFeatures<?> features) {
MockKey<?> mockKey = MockKey.of(features.mockedType, features.interfaces);
Class<?> generatedMockClass = getMockClass(mockKey);
if (generatedMockClass == null) {
synchronized (features.mockedType) {
generatedMockClass = getMockClass(mockKey);
if(generatedMockClass == null) {
generatedMockClass = generate(features);
generatedClassCache.put(mockKey, weak ? new WeakReference<Class<?>>(generatedMockClass) : new SoftReference<Class<?>>(generatedMockClass));
}
}
}
generatedClassCache.put(mockKey, new WeakReference<Class<?>>(generatedMockClass));
return generatedMockClass;
}

Expand Down Expand Up @@ -98,14 +123,19 @@ private RuntimeException prettifyFailure(MockFeatures<?> mockFeatures, Exception
// should be stored as a weak reference
private static class MockKey<T> {
private final String mockedType;
private final Set<String> types = new HashSet<String>();
private final Set<String> types;

private MockKey(Class<T> mockedType, Set<Class<?>> interfaces) {
this.mockedType = mockedType.getName();
for (Class<?> anInterface : interfaces) {
types.add(anInterface.getName());
if (interfaces.isEmpty()) { // Optimize memory footprint for the common case.
types = Collections.emptySet();
} else {
types = new HashSet<String>();
for (Class<?> anInterface : interfaces) {
types.add(anInterface.getName());
}
types.add(this.mockedType);
}
types.add(this.mockedType);
}

@Override
Expand Down Expand Up @@ -133,4 +163,60 @@ public static <T> MockKey<T> of(Class<T> mockedType, Set<Class<?>> interfaces) {
}
}
}

private interface Key<T> {

T get();
}

private static class StrongKey<T> implements Key<T> {

private final T value;

private final int hashCode;

public StrongKey(T value) {
this.value = value;
hashCode = System.identityHashCode(value);
}

@Override
public T get() {
return value;
}

@Override
public boolean equals(Object object) {
if (this == object) return true;
if (!(object instanceof Key)) return false;
return value == ((Key<?>) object).get();
}

@Override
public int hashCode() {
return hashCode;
}
}

private static class WeakKey<T> extends WeakReference<T> implements Key<T> {

private final int hashCode;

public WeakKey(T referent, ReferenceQueue<? super T> q) {
super(referent, q);
hashCode = System.identityHashCode(referent);
}

@Override
public boolean equals(Object object) {
if (this == object) return true;
if (!(object instanceof Key)) return false;
return get() == ((Key<?>) object).get();
}

@Override
public int hashCode() {
return hashCode;
}
}
}
5 changes: 4 additions & 1 deletion src/test/java/org/mockito/internal/creation/bytebuddy/CachingMockBytecodeGeneratorTest.java
View file
@@ -1,6 +1,7 @@
package org.mockito.internal.creation.bytebuddy;

import org.junit.Before;
import org.junit.Ignore;
import org.junit.Test;

import java.lang.management.ManagementFactory;
Expand Down Expand Up @@ -29,7 +30,7 @@ public void ensure_cache_is_cleared_if_no_reference_to_classloader_and_classes()
.withClassDefinition("foo.Bar", makeMarkerInterface("foo.Bar"))
.build();

CachingMockBytecodeGenerator cachingMockBytecodeGenerator = new CachingMockBytecodeGenerator();
CachingMockBytecodeGenerator cachingMockBytecodeGenerator = new CachingMockBytecodeGenerator(true);
Class<?> the_mock_type = cachingMockBytecodeGenerator.get(withMockFeatures(
classloader_with_life_shorter_than_cache.loadClass("foo.Bar"),
Collections.<Class<?>>emptySet(),
Expand All @@ -45,6 +46,8 @@ public void ensure_cache_is_cleared_if_no_reference_to_classloader_and_classes()
System.gc();
ensure_gc_happened();

cachingMockBytecodeGenerator.pollAndClearReferences();

// then
assertThat(cachingMockBytecodeGenerator.avoidingClassLeakageCache).isEmpty();
}
Expand Down

0 comments on commit 0ff47d4

Please sign in to comment.