Fix race condition in inline static mocking

dexmaker-mockito-inline-extended-tests/src/androidTest/java/com/android/dx/mockito/inline/extended/tests/ConcurrentStaticMocking.java

@@ -0,0 +1,106 @@
+/*
+ * Copyright (C) 2026 The Android Open Source Project
+ *
+ * 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
+ *
+ *      http://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 com.android.dx.mockito.inline.extended.tests;
+
+import org.junit.Test;
+import org.mockito.MockitoSession;
+
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
+
+import static com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession;
+import static com.android.dx.mockito.inline.extended.ExtendedMockito.staticMockMarker;
+import static org.junit.Assert.assertNull;
+
+public class ConcurrentStaticMocking {
+
+    public static class SuperClass {
+        public static void unmockedMethod() {}
+    }
+
+    public static class SubClass extends SuperClass {
+    }
+
+    private static class DummyClass {
+        static void dummyMethod() {}
+    }
+
+    @Test
+    public void concurrentMockingWithBackgroundThread() throws Exception {
+        AtomicBoolean running = new AtomicBoolean(true);
+        AtomicReference<Throwable> bgError = new AtomicReference<>();
+        CountDownLatch threadStarted = new CountDownLatch(1);
+
+        // 1. One-time setup: Spying SubClass forces Dexmaker to instrument SuperClass
+        MockitoSession initSession = mockitoSession()
+                .spyStatic(SubClass.class)
+                .startMocking();
+        initSession.finishMocking();
+
+        // 2. Start background thread continuously executing unmocked method on SuperClass.
+        // Dexmaker intercepts this call to check if any mocked subclass handles it,
+        // which iterates over classToMarker.keySet().
+        Thread bgThread = new Thread(() -> {
+            threadStarted.countDown();
+            while (running.get()) {
+                try {
+                    SuperClass.unmockedMethod();
+                } catch (Throwable t) {
+                    bgError.set(t);
+                    break;
+                }
+            }
+        }, "ConcurrentTestWorker");
+        bgThread.start();
+        threadStarted.await();
+
+        // 3. Concurrently cycle through creating and finishing static mocking sessions,
+        // which mutates classToMarker concurrently with the background thread's iteration.
+        try {
+            for (int i = 0; i < 500; i++) {
+                if (bgError.get() != null) {
+                    break;
+                }
+                MockitoSession session = mockitoSession()
+                        .spyStatic(SubClass.class)
+                        .mockStatic(DummyClass.class)
+                        .startMocking();
+                session.finishMocking();
+            }
+        } finally {
+            running.set(false);
+            bgThread.join(5000);
+        }
+
+        assertNull("Background thread threw exception: " + bgError.get(), bgError.get());
+    }
+
+    @Test
+    public void staticMockMarkerWithNullClass() {
+        assertNull(staticMockMarker((Class<?>) null));
+
+        MockitoSession session = mockitoSession()
+                .mockStatic(DummyClass.class)
+                .startMocking();
+        try {
+            assertNull(staticMockMarker((Class<?>) null));
+        } finally {
+            session.finishMocking();
+        }
+    }
+}

dexmaker-mockito-inline-extended/src/main/java/com/android/dx/mockito/inline/InlineStaticMockMaker.java

@@ -30,9 +30,9 @@
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
 import java.lang.reflect.Modifier;
-import java.util.HashMap;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.function.BiConsumer;
 
 /**
@@ -113,7 +113,10 @@ public final class InlineStaticMockMaker implements MockMaker {
      * object's method calls should be intercepted.
      */
     private final Map<Object, InvocationHandlerAdapter> markerToHandler = new MarkerToHandlerMap();
-    private final Map<Class, Object> classToMarker = new HashMap<>();
+    // NOTE: Must be ConcurrentHashMap to prevent ConcurrentModificationException when
+    // background threads execute methods on instrumented classes while test threads mutate this map.
+    // Also note that ConcurrentHashMap does not accept {@code null} keys or values.
+    private final Map<Class, Object> classToMarker = new ConcurrentHashMap<>();
 
     /**
      * Class doing the actual byte code transformation.

dexmaker-mockito-inline-extended/src/main/java/com/android/dx/mockito/inline/MarkerToHandlerMap.java

@@ -5,10 +5,10 @@
 
 import java.util.AbstractMap;
 import java.util.Collection;
-import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
 
 /**
  * A map for mock marker object -> {@link InvocationHandlerAdapter} but
@@ -20,7 +20,11 @@
  */
 class MarkerToHandlerMap implements Map<Object, InvocationHandlerAdapter> {
 
-    private final Map<MockMarkerKey, InvocationHandlerAdapter> markerToHandler = new HashMap<>();
+    // NOTE: Must be ConcurrentHashMap to ensure thread safety when background threads
+    // perform mock dispatch while test threads register or remove handlers. Note that
+    // ConcurrentHashMap does not accept {@code null} keys or values, which is why null checks
+    // are explicitly handled in this map.
+    private final Map<MockMarkerKey, InvocationHandlerAdapter> markerToHandler = new ConcurrentHashMap<>();
 
     @Override
     public int size() {
@@ -34,26 +38,41 @@ public boolean isEmpty() {
 
     @Override
     public boolean containsKey(Object key) {
+        if (key == null) {
+            return false;
+        }
         return markerToHandler.containsKey(new MockMarkerKey(key));
     }
 
     @Override
     public boolean containsValue(Object value) {
+        if (value == null) {
+            return false;
+        }
         return markerToHandler.containsValue(value);
     }
 
     @Override
     public InvocationHandlerAdapter get(Object key) {
+        if (key == null) {
+            return null;
+        }
         return markerToHandler.get(new MockMarkerKey(key));
     }
 
     @Override
     public InvocationHandlerAdapter put(Object key, InvocationHandlerAdapter value) {
+        if (key == null || value == null) {
+            return null;
+        }
         return markerToHandler.put(new MockMarkerKey(key), value);
     }
 
     @Override
     public InvocationHandlerAdapter remove(Object key) {
+        if (key == null) {
+            return null;
+        }
         return markerToHandler.remove(new MockMarkerKey(key));
     }
 
@@ -71,7 +90,7 @@ public void clear() {
 
     @Override
     public Set<Object> keySet() {
-        Set<Object> set = new HashSet<>(entrySet().size());
+        Set<Object> set = new HashSet<>(markerToHandler.size());
         for (MockMarkerKey key : markerToHandler.keySet()) {
             set.add(key.mockMarker);
         }
@@ -86,7 +105,7 @@ public Collection<InvocationHandlerAdapter> values() {
     @Override
     @SuppressWarnings("InfiniteRecursion")
     public Set<Entry<Object, InvocationHandlerAdapter>> entrySet() {
-        Set<Entry<Object, InvocationHandlerAdapter>> set = new HashSet<>(entrySet().size());
+        Set<Entry<Object, InvocationHandlerAdapter>> set = new HashSet<>(markerToHandler.size());
         for (Entry<MockMarkerKey, InvocationHandlerAdapter> entry : markerToHandler.entrySet()) {
             set.add(new AbstractMap.SimpleImmutableEntry<>(entry.getKey().mockMarker, entry.getValue()));
         }

dexmaker-mockito-inline-extended/src/main/java/com/android/dx/mockito/inline/StaticMockMethodAdvice.java

@@ -227,8 +227,14 @@ public Callable<?> handle(Object methodDescStr, Method origin, Object[] argument
             Throwable {
         MethodDesc methodDesc = new MethodDesc((String) methodDescStr);
         Class clazz = getClassMethodWasCalledOn(methodDesc);
+        if (clazz == null) {
+            return null;
+        }
 
         Object marker = classToMarker.get(clazz);
+        if (marker == null) {
+            return null;
+        }
         InvocationHandlerAdapter interceptor = markersToHandler.get(marker);
         if (interceptor == null) {
             return null;
@@ -256,6 +262,9 @@ public Callable<?> handle(Object methodDescStr, Method origin, Object[] argument
      * @return {@code true} iff the marker is a mock marker
      */
     public boolean isMarker(Object marker) {
+        if (marker == null) {
+            return false;
+        }
         return markersToHandler.containsKey(marker);
     }
 

dexmaker-mockito-inline-extended/src/main/java/com/android/dx/mockito/inline/extended/ExtendedMockito.java

@@ -29,6 +29,7 @@
 import java.lang.reflect.Method;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.concurrent.CopyOnWriteArrayList;
 
 import static com.android.dx.mockito.inline.InlineDexmakerMockMaker.onSpyInProgressInstance;
 import static com.android.dx.mockito.inline.InlineStaticMockMaker.onMethodCallDuringVerification;
@@ -66,8 +67,12 @@
 public class ExtendedMockito extends Mockito {
     /**
      * Currently active {@link #mockitoSession() sessions}
+     *
+     * <p>NOTE: Must be {@link CopyOnWriteArrayList} to prevent {@link java.util
+     * .ConcurrentModificationException} when background threads query active sessions
+     * while test threads start or finish mocking sessions.
      */
-    private static ArrayList<StaticMockitoSession> sessions = new ArrayList<>();
+    private static final List<StaticMockitoSession> sessions = new CopyOnWriteArrayList<>();
 
     /**
      * Same as {@link Mockito#doAnswer(Answer)} but adds the ability to stub static method calls via

dexmaker-mockito-inline-extended/src/main/java/com/android/dx/mockito/inline/extended/StaticMockitoSession.java

@@ -21,7 +21,8 @@
 import org.mockito.quality.Strictness;
 
 import java.util.ArrayList;
-import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 
 import static com.android.dx.mockito.inline.InlineStaticMockMaker.mockingInProgressClass;
 
@@ -32,8 +33,13 @@
 public class StaticMockitoSession implements MockitoSession {
     /**
      * For each class where static mocking is enabled there is one marker object.
+     *
+     * <p>NOTE: Must be {@link ConcurrentHashMap} to prevent {@link java.util
+     * .ConcurrentModificationException} when background threads execute methods on
+     * instrumented classes while test threads mutate this map. Note that
+     * {@link ConcurrentHashMap} does not accept {@code null} keys or values.
      */
-    private static final HashMap<Class, Object> classToMarker = new HashMap<>();
+    private static final Map<Class, Object> classToMarker = new ConcurrentHashMap<>();
 
     private final MockitoSession instanceSession;
     private final ArrayList<Class<?>> staticMocks = new ArrayList<>(0);
@@ -112,6 +118,9 @@ <T> void mockStatic(StaticMocking<T> mocking) {
      */
     @SuppressWarnings("unchecked")
     <T> T staticMockMarker(Class<T> clazz) {
+        if (clazz == null) {
+            return null;
+        }
         return (T) classToMarker.get(clazz);
     }
 }