Use multiple locks in ByteBuddyMockFactory to reduce lock contention (#1778)

AndreasTu · web-flow · commit d660020f49d5 · 2023-09-11T00:14:14.000+02:00
Changed locking scheme in ByteBuddyMockFactory from a single global lock CACHE to cacheLocks with size 16. The used TypeCachingLock from cacheLocks depends on the hashcode of TypeCache.SimpleKey, which aggregates the types to mock. The old global CACHE lock did block Threads regardless, if they try to mock the same type or not. This is a similar fix as in Mockito: mockito/mockito#3095
diff --git a/spock-core/src/main/java/org/spockframework/mock/runtime/ByteBuddyMockFactory.java b/spock-core/src/main/java/org/spockframework/mock/runtime/ByteBuddyMockFactory.java
@@ -1,3 +1,19 @@
+/*
+ * Copyright 2023 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
 package org.spockframework.mock.runtime;
 
 import net.bytebuddy.ByteBuddy;
@@ -13,33 +29,66 @@
 import net.bytebuddy.implementation.MethodDelegation;
 import net.bytebuddy.implementation.bind.annotation.Morph;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.VisibleForTesting;
 import org.spockframework.mock.ISpockMockObject;
 import org.spockframework.mock.codegen.Target;
 import org.spockframework.util.Nullable;
 
 import java.lang.reflect.Method;
 import java.util.List;
+import java.util.concurrent.Callable;
 import java.util.concurrent.ThreadLocalRandom;
 
 import static net.bytebuddy.matcher.ElementMatchers.any;
 import static net.bytebuddy.matcher.ElementMatchers.none;
 
 class ByteBuddyMockFactory {
+  /**
+   * The mask to use to mask out the {@link TypeCache.SimpleKey#hashCode()} to find the {@link #cacheLocks}.
+   */
+  private static final int CACHE_LOCK_MASK = 0x0F;
+
+  /**
+   * The size of the {@link #cacheLocks}.
+   */
+  private static final int CACHE_LOCK_SIZE = CACHE_LOCK_MASK + 1;
 
-  private static final TypeCache<TypeCache.SimpleKey> CACHE =
-    new TypeCache.WithInlineExpunction<>(TypeCache.Sort.SOFT);
   private static final Class<?> CODEGEN_TARGET_CLASS = Target.class;
   private static final String CODEGEN_PACKAGE = CODEGEN_TARGET_CLASS.getPackage().getName();
 
-  static Object createMock(final Class<?> type,
+  /**
+   * This array contains {@link TypeCachingLock} instances, which are used as java monitor locks for
+   * {@link TypeCache#findOrInsert(ClassLoader, Object, Callable, Object)}.
+   * The {@code cacheLocks} spreads the lock to acquire over multiple locks instead of using a single lock
+   * {@code CACHE} for all {@link TypeCache.SimpleKey}s.
+   *
+   * <p>Note: We can't simply use the mockedType class lock as a lock,
+   * because the {@code TypeCache.SimpleKey}, will be the same for different {@code mockTypes + additionalInterfaces}.
+   * See the {@code hashCode} implementation of the {@code TypeCache.SimpleKey}, which has {@code Set} semantics.
+   */
+  private final TypeCachingLock[] cacheLocks;
+
+  private final TypeCache<TypeCache.SimpleKey> CACHE =
+    new TypeCache.WithInlineExpunction<>(TypeCache.Sort.SOFT);
+
+  ByteBuddyMockFactory() {
+    cacheLocks = new TypeCachingLock[CACHE_LOCK_SIZE];
+    for (int i = 0; i < CACHE_LOCK_SIZE; i++) {
+      cacheLocks[i] = new TypeCachingLock();
+    }
+  }
+
+  Object createMock(final Class<?> type,
                            final List<Class<?>> additionalInterfaces,
                            @Nullable List<Object> constructorArgs,
                            IProxyBasedMockInterceptor interceptor,
                            final ClassLoader classLoader,
                            boolean useObjenesis) {
 
+    TypeCache.SimpleKey key = new TypeCache.SimpleKey(type, additionalInterfaces);
+
     Class<?> enhancedType = CACHE.findOrInsert(classLoader,
-      new TypeCache.SimpleKey(type, additionalInterfaces),
+      key,
       () -> {
         String typeName = type.getName();
         Class<?> targetClass = type;
@@ -68,14 +117,28 @@ static Object createMock(final Class<?> type,
           .make()
           .load(classLoader, strategy)
           .getLoaded();
-      }, CACHE);
+      }, getCacheLockForKey(key));
 
     Object proxy = MockInstantiator.instantiate(type, enhancedType, constructorArgs, useObjenesis);
     ((ByteBuddyInterceptorAdapter.InterceptorAccess) proxy).$spock_set(interceptor);
     return proxy;
   }
 
-  // This methods and the ones it calls are inspired by similar code in Mockito's SubclassBytecodeGenerator
+  /**
+   * Returns a {@link TypeCachingLock}, which locks the {@link TypeCache#findOrInsert(ClassLoader, Object, Callable, Object)}.
+   *
+   * @param key the key to lock
+   * @return the {@link TypeCachingLock} to use to lock the {@link TypeCache}
+   */
+  private TypeCachingLock getCacheLockForKey(TypeCache.SimpleKey key) {
+    int hashCode = key.hashCode();
+    // Try to spread some higher bits with XOR to lower bits, because we only use lower bits.
+    hashCode = hashCode ^ (hashCode >>> 16);
+    int index = hashCode & CACHE_LOCK_MASK;
+    return cacheLocks[index];
+  }
+
+  // This method and the ones it calls are inspired by similar code in Mockito's SubclassBytecodeGenerator
   private static boolean shouldLoadIntoCodegenPackage(Class<?> type) {
     return isComingFromJDK(type) || isComingFromSignedJar(type) || isComingFromSealedPackage(type);
   }
@@ -112,4 +175,6 @@ private static ClassLoadingStrategy<ClassLoader> determineBestClassLoadingStrate
     return ClassLoadingStrategy.Default.WRAPPER;
   }
 
+  private static final class TypeCachingLock {
+  }
 }
diff --git a/spock-core/src/main/java/org/spockframework/mock/runtime/ProxyBasedMockFactory.java b/spock-core/src/main/java/org/spockframework/mock/runtime/ProxyBasedMockFactory.java
@@ -33,13 +33,24 @@ public class ProxyBasedMockFactory {
 
   public static ProxyBasedMockFactory INSTANCE = new ProxyBasedMockFactory();
 
+  @Nullable
+  private final ByteBuddyMockFactory byteBuddyMockFactory;
+
+  private ProxyBasedMockFactory() {
+    if (byteBuddyAvailable && !ignoreByteBuddy) {
+      byteBuddyMockFactory = new ByteBuddyMockFactory();
+    } else {
+      byteBuddyMockFactory = null;
+    }
+  }
+
   public Object create(Class<?> mockType, List<Class<?>> additionalInterfaces, @Nullable List<Object> constructorArgs,
       IProxyBasedMockInterceptor mockInterceptor, ClassLoader classLoader, boolean useObjenesis) throws CannotCreateMockException {
     if (mockType.isInterface()) {
       return DynamicProxyMockFactory.createMock(mockType, additionalInterfaces, constructorArgs, mockInterceptor, classLoader);
     }
-    if (byteBuddyAvailable && !ignoreByteBuddy) {
-      return ByteBuddyMockFactory.createMock(mockType, additionalInterfaces, constructorArgs, mockInterceptor, classLoader, useObjenesis);
+    if (byteBuddyMockFactory != null) {
+      return byteBuddyMockFactory.createMock(mockType, additionalInterfaces, constructorArgs, mockInterceptor, classLoader, useObjenesis);
     }
     if (cglibAvailable) {
       return CglibMockFactory.createMock(mockType, additionalInterfaces, constructorArgs, mockInterceptor, classLoader, useObjenesis);
diff --git a/spock-specs/src/test/groovy/org/spockframework/mock/runtime/ByteBuddyMockFactoryConcurrentSpec.groovy b/spock-specs/src/test/groovy/org/spockframework/mock/runtime/ByteBuddyMockFactoryConcurrentSpec.groovy
@@ -0,0 +1,146 @@
+/*
+ * Copyright 2023 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.spockframework.mock.runtime
+
+import groovy.transform.Canonical
+import spock.lang.Shared
+import spock.lang.Specification
+import spock.lang.Timeout
+
+import java.util.concurrent.CompletableFuture
+import java.util.concurrent.Phaser
+import java.util.concurrent.TimeUnit
+import java.util.function.Function
+
+class ByteBuddyMockFactoryConcurrentSpec extends Specification {
+  private static final String IfA = "IfA"
+  private static final String IfB = "IfB"
+  private static final String IfC = "IfC"
+  @Shared
+  final IProxyBasedMockInterceptor interceptor = Stub()
+
+  def "ensure lockMask bit patterns"() {
+    expect:
+    1 << Integer.bitCount(ByteBuddyMockFactory.CACHE_LOCK_MASK) - 1 == Integer.highestOneBit(ByteBuddyMockFactory.CACHE_LOCK_MASK)
+  }
+
+  // Just to be save to abort, normally the tests run in 2 secs.
+  @Timeout(40)
+  def "cacheLockingStressTest #test"() {
+    given:
+    int iterations = 5_000
+    def tempClassLoader = new ByteBuddyTestClassLoader()
+    MockFeatures featA = toMockFeatures(mockSpecA, tempClassLoader)
+    MockFeatures featB = toMockFeatures(mockSpecB, tempClassLoader)
+    ByteBuddyMockFactory mockFactory = new ByteBuddyMockFactory()
+
+    Phaser phaser = new Phaser(4)
+    Function<Runnable, CompletableFuture<Void>> runCode = { Runnable code ->
+      CompletableFuture.runAsync {
+        phaser.arriveAndAwaitAdvance()
+        try {
+          for (int i = 0; i < iterations; i++) {
+            code.run()
+          }
+        } finally {
+          phaser.arrive()
+        }
+      }
+    }
+    when:
+    def mockFeatAFuture = runCode.apply {
+      Class<?> mockClass = mockClass(mockFactory, tempClassLoader, featA)
+      assertValidMockClass(featA, mockClass, tempClassLoader)
+    }
+
+    def mockFeatBFuture = runCode.apply {
+      Class<?> mockClass = mockClass(mockFactory, tempClassLoader, featB)
+      assertValidMockClass(featB, mockClass, tempClassLoader)
+    }
+
+    def cacheFuture = runCode.apply { mockFactory.CACHE.clear() }
+
+    phaser.arriveAndAwaitAdvance()
+    // Wait for test to end
+    int phase = phaser.arrive()
+    try {
+      phaser.awaitAdvanceInterruptibly(phase, 30, TimeUnit.SECONDS)
+    } finally {
+      // Collect exceptions from the futures, to make issues visible.
+      mockFeatAFuture.getNow(null)
+      mockFeatBFuture.getNow(null)
+      cacheFuture.getNow(null)
+    }
+    then:
+    noExceptionThrown()
+
+    where:
+    test                                      | mockSpecA          | mockSpecB
+    "same hashcode different mockType"        | mockSpec(IfA, IfB) | mockSpec(IfB, IfA)
+    "same hashcode same mockType"             | mockSpec(IfA)      | mockSpec(IfA)
+    "different hashcode different interfaces" | mockSpec(IfA, IfB) | mockSpec(IfB, IfC)
+    "unrelated classes"                       | mockSpec(IfA)      | mockSpec(IfB)
+  }
+
+  private Class<?> mockClass(ByteBuddyMockFactory mockFactory, ClassLoader cl, MockFeatures feature) {
+    return mockFactory.createMock(
+      feature.mockType,
+      feature.interfaces,
+      null,
+      interceptor,
+      cl,
+      false)
+      .getClass()
+  }
+
+  private static MockSpec mockSpec(String mockedType, String... interfaces) {
+    return new MockSpec(mockedType, interfaces as List)
+  }
+
+  private void assertValidMockClass(MockFeatures mockFeature, Class<?> mockClass, ClassLoader classLoader) {
+    assert mockClass.classLoader == classLoader
+    assert mockFeature.mockType.isAssignableFrom(mockClass)
+    mockFeature.interfaces.each {
+      assert it.isAssignableFrom(mockClass)
+    }
+  }
+
+  MockFeatures toMockFeatures(MockSpec mockFeaturesString, ByteBuddyTestClassLoader classLoader) {
+    def mockType = classLoader.defineInterface(mockFeaturesString.mockType)
+    def interfaces = mockFeaturesString.interfaces.collect { classLoader.defineInterface(it) }
+    return new MockFeatures(mockType, interfaces)
+  }
+
+  /**
+   * Class holding the loaded classes to mock.
+   */
+  @Canonical
+  private static class MockFeatures {
+    final Class<?> mockType
+    final List<Class<?>> interfaces
+  }
+
+  /**
+   * Class holding the class names to mock.
+   * Which will be converted into a {@link MockFeatures} during test.
+   */
+  @Canonical
+  private static class MockSpec {
+    final String mockType
+    final List<String> interfaces
+  }
+}
diff --git a/spock-specs/src/test/groovy/org/spockframework/mock/runtime/ByteBuddyTestClassLoader.java b/spock-specs/src/test/groovy/org/spockframework/mock/runtime/ByteBuddyTestClassLoader.java
@@ -0,0 +1,55 @@
+/*
+ * Copyright 2023 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.spockframework.mock.runtime;
+
+import net.bytebuddy.ByteBuddy;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.dynamic.loading.ClassLoadingStrategy;
+
+import java.util.HashMap;
+import java.util.Map;
+
+final class ByteBuddyTestClassLoader extends ClassLoader {
+  private final Map<String, Class<?>> cache = new HashMap<>();
+
+  private final ClassLoadingStrategy<ByteBuddyTestClassLoader> loadingStrategy = (loader, types) -> {
+    Map<TypeDescription, Class<?>> result = new HashMap<>();
+    for (Map.Entry<TypeDescription, byte[]> entry : types.entrySet()) {
+      TypeDescription description = entry.getKey();
+      byte[] bytes = entry.getValue();
+      Class<?> clazz = defineClass(description.getName(), bytes, 0, bytes.length);
+      result.put(description, clazz);
+    }
+    return result;
+  };
+
+  /**
+   * Defines an empty interface with the passed {@code node}.
+   *
+   * @param name the name of the interface
+   * @return the loaded {@code Class}
+   */
+  synchronized Class<?> defineInterface(String name) {
+    //noinspection resource
+    return cache.computeIfAbsent(name, nameKey -> new ByteBuddy()
+      .makeInterface()
+      .name(nameKey)
+      .make()
+      .load(this, loadingStrategy)
+      .getLoaded());
+  }
+}
