Use multiple locks in ByteBuddyMockFactory to reduce lock contention

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 {
+  }
 }

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);

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
+  }
+}

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());
+  }
+}