diff --git a/javalib/src/android/system/virtualmachine/VirtualMachine.java b/javalib/src/android/system/virtualmachine/VirtualMachine.java
index c200d00..d9c75c0 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachine.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachine.java
@@ -76,7 +76,6 @@
 import java.io.InputStreamReader;
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
-import java.lang.ref.WeakReference;
 import java.nio.channels.FileChannel;
 import java.nio.file.FileAlreadyExistsException;
 import java.nio.file.FileVisitResult;
@@ -86,10 +85,7 @@
 import java.nio.file.attribute.BasicFileAttributes;
 import java.util.ArrayList;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.List;
-import java.util.Map;
-import java.util.WeakHashMap;
 import java.util.concurrent.Executor;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Consumer;
@@ -107,11 +103,6 @@
  * @hide
  */
 public class VirtualMachine implements AutoCloseable {
-    /** Map from context to a map of all that context's VMs by name. */
-    @GuardedBy("sCreateLock")
-    private static final Map<Context, Map<String, WeakReference<VirtualMachine>>> sInstances =
-            new WeakHashMap<>();
-
     /** Name of the directory under the files directory where all VMs created for the app exist. */
     private static final String VM_DIR = "vm";
 
@@ -208,17 +199,10 @@
     private static final long INSTANCE_FILE_SIZE = 10 * 1024 * 1024;
 
     // A note on lock ordering:
-    // You can take mLock while holding sCreateLock, but not vice versa.
+    // You can take mLock while holding VirtualMachineManager.sCreateLock, but not vice versa.
     // We never take any other lock while holding mCallbackLock; therefore you can
     // take mCallbackLock while holding any other lock.
 
-    /**
-     * A lock used to synchronize the creation of virtual machines. It protects
-     * {@link #sInstances}, but is also held throughout VM creation / retrieval / deletion, to
-     * prevent these actions racing with each other.
-     */
-    static final Object sCreateLock = new Object();
-
     /** Lock protecting our mutable state (other than callbacks). */
     private final Object mLock = new Object();
 
@@ -281,12 +265,6 @@
         mExtraApks = setupExtraApks(context, config, thisVmDir);
     }
 
-    @GuardedBy("sCreateLock")
-    @NonNull
-    private static Map<String, WeakReference<VirtualMachine>> getInstancesMap(Context context) {
-        return sInstances.computeIfAbsent(context, unused -> new HashMap<>());
-    }
-
     /**
      * Builds a virtual machine from an {@link VirtualMachineDescriptor} object and associates it
      * with the given name.
@@ -297,7 +275,7 @@
      * #delete}. The imported virtual machine is in {@link #STATUS_STOPPED} state. To run the VM,
      * call {@link #run}.
      */
-    @GuardedBy("sCreateLock")
+    @GuardedBy("VirtualMachineManager.sCreateLock")
     @NonNull
     static VirtualMachine fromDescriptor(
             @NonNull Context context,
@@ -315,7 +293,6 @@
                 throw new VirtualMachineException("failed to create instance image", e);
             }
             vm.importInstanceFrom(vmDescriptor.getInstanceImgFd());
-            getInstancesMap(context).put(name, new WeakReference<>(vm));
             return vm;
         } catch (VirtualMachineException | RuntimeException e) {
             // If anything goes wrong, delete any files created so far and the VM's directory
@@ -333,7 +310,7 @@
      * it is persisted until it is deleted by calling {@link #delete}. The created virtual machine
      * is in {@link #STATUS_STOPPED} state. To run the VM, call {@link #run}.
      */
-    @GuardedBy("sCreateLock")
+    @GuardedBy("VirtualMachineManager.sCreateLock")
     @NonNull
     static VirtualMachine create(
             @NonNull Context context, @NonNull String name, @NonNull VirtualMachineConfig config)
@@ -365,9 +342,6 @@
             } catch (ServiceSpecificException | IllegalArgumentException e) {
                 throw new VirtualMachineException("failed to create instance partition", e);
             }
-
-            getInstancesMap(context).put(name, new WeakReference<>(vm));
-
             return vm;
         } catch (VirtualMachineException | RuntimeException e) {
             // If anything goes wrong, delete any files created so far and the VM's directory
@@ -381,10 +355,10 @@
     }
 
     /** Loads a virtual machine that is already created before. */
-    @GuardedBy("sCreateLock")
+    @GuardedBy("VirtualMachineManager.sCreateLock")
     @Nullable
-    static VirtualMachine load(
-            @NonNull Context context, @NonNull String name) throws VirtualMachineException {
+    static VirtualMachine load(@NonNull Context context, @NonNull String name)
+            throws VirtualMachineException {
         File thisVmDir = getVmDir(context, name);
         if (!thisVmDir.exists()) {
             // The VM doesn't exist.
@@ -392,51 +366,33 @@
         }
         File configFilePath = new File(thisVmDir, CONFIG_FILE);
         VirtualMachineConfig config = VirtualMachineConfig.from(configFilePath);
-        Map<String, WeakReference<VirtualMachine>> instancesMap = getInstancesMap(context);
-
-        VirtualMachine vm = null;
-        if (instancesMap.containsKey(name)) {
-            vm = instancesMap.get(name).get();
-        }
-        if (vm == null) {
-            vm = new VirtualMachine(context, name, config);
-        }
+        VirtualMachine vm = new VirtualMachine(context, name, config);
 
         if (!vm.mInstanceFilePath.exists()) {
             throw new VirtualMachineException("instance image missing");
         }
 
-        instancesMap.put(name, new WeakReference<>(vm));
-
         return vm;
     }
 
-    @GuardedBy("sCreateLock")
-    static void delete(Context context, String name) throws VirtualMachineException {
-        Map<String, WeakReference<VirtualMachine>> instancesMap = sInstances.get(context);
-        VirtualMachine vm;
-        if (instancesMap != null && instancesMap.containsKey(name)) {
-            vm = instancesMap.get(name).get();
-        } else {
-            vm = null;
+    @GuardedBy("VirtualMachineManager.sCreateLock")
+    void delete(Context context, String name) throws VirtualMachineException {
+        synchronized (mLock) {
+            checkStopped();
         }
 
-        if (vm != null) {
-            synchronized (vm.mLock) {
-                vm.checkStopped();
-            }
-        }
+        deleteVmDirectory(context, name);
+    }
 
+    static void deleteVmDirectory(Context context, String name) throws VirtualMachineException {
         try {
             deleteRecursively(getVmDir(context, name));
         } catch (IOException e) {
             throw new VirtualMachineException(e);
         }
-
-        if (instancesMap != null) instancesMap.remove(name);
     }
 
-    @GuardedBy("sCreateLock")
+    @GuardedBy("VirtualMachineManager.sCreateLock")
     @NonNull
     private static File createVmDir(@NonNull Context context, @NonNull String name)
             throws VirtualMachineException {
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachineManager.java b/javalib/src/android/system/virtualmachine/VirtualMachineManager.java
index c357f50..098e3ca 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachineManager.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachineManager.java
@@ -25,6 +25,7 @@
 import android.annotation.SuppressLint;
 import android.content.Context;
 import android.sysprop.HypervisorProperties;
+import android.util.ArrayMap;
 
 import com.android.internal.annotations.GuardedBy;
 
@@ -47,6 +48,13 @@
  * @hide
  */
 public class VirtualMachineManager {
+    /**
+     * A lock used to synchronize the creation of virtual machines. It protects {@link #sInstances},
+     * but is also held throughout VM creation / retrieval / deletion, to prevent these actions
+     * racing with each other.
+     */
+    private static final Object sCreateLock = new Object();
+
     @NonNull private final Context mContext;
 
     private VirtualMachineManager(@NonNull Context context) {
@@ -57,6 +65,10 @@
     private static final Map<Context, WeakReference<VirtualMachineManager>> sInstances =
             new WeakHashMap<>();
 
+    @NonNull
+    @GuardedBy("sCreateLock")
+    private final Map<String, WeakReference<VirtualMachine>> mVmsByName = new ArrayMap<>();
+
     /**
      * Capabilities of the virtual machine implementation.
      *
@@ -136,11 +148,48 @@
     public VirtualMachine create(
             @NonNull String name, @NonNull VirtualMachineConfig config)
             throws VirtualMachineException {
-        synchronized (VirtualMachine.sCreateLock) {
-            return VirtualMachine.create(mContext, name, config);
+        synchronized (sCreateLock) {
+            return createLocked(name, config);
         }
     }
 
+    @NonNull
+    @GuardedBy("sCreateLock")
+    private VirtualMachine createLocked(String name, VirtualMachineConfig config)
+            throws VirtualMachineException {
+        VirtualMachine vm = VirtualMachine.create(mContext, name, config);
+        mVmsByName.put(name, new WeakReference<>(vm));
+        return vm;
+    }
+
+    /**
+     * Returns an existing {@link VirtualMachine} with the given name. Returns null if there is no
+     * such virtual machine.
+     *
+     * @throws VirtualMachineException if the virtual machine exists but could not be successfully
+     *     retrieved.
+     * @hide
+     */
+    @Nullable
+    public VirtualMachine get(@NonNull String name) throws VirtualMachineException {
+        synchronized (sCreateLock) {
+            return getLocked(name);
+        }
+    }
+
+    @Nullable
+    @GuardedBy("sCreateLock")
+    private VirtualMachine getLocked(String name) throws VirtualMachineException {
+        VirtualMachine vm = getVmByName(name);
+        if (vm != null) return vm;
+
+        vm = VirtualMachine.load(mContext, name);
+        if (vm != null) {
+            mVmsByName.put(name, new WeakReference<>(vm));
+        }
+        return vm;
+    }
+
     /**
      * Imports a virtual machine from an {@link VirtualMachineDescriptor} object and associates it
      * with the given name.
@@ -154,23 +203,10 @@
     public VirtualMachine importFromDescriptor(
             @NonNull String name, @NonNull VirtualMachineDescriptor vmDescriptor)
             throws VirtualMachineException {
-        synchronized (VirtualMachine.sCreateLock) {
-            return VirtualMachine.fromDescriptor(mContext, name, vmDescriptor);
-        }
-    }
-
-    /**
-     * Returns an existing {@link VirtualMachine} with the given name. Returns null if there is no
-     * such virtual machine.
-     *
-     * @throws VirtualMachineException if the virtual machine exists but could not be successfully
-     *                                 retrieved.
-     * @hide
-     */
-    @Nullable
-    public VirtualMachine get(@NonNull String name) throws VirtualMachineException {
-        synchronized (VirtualMachine.sCreateLock) {
-            return VirtualMachine.load(mContext, name);
+        synchronized (sCreateLock) {
+            VirtualMachine vm = VirtualMachine.fromDescriptor(mContext, name, vmDescriptor);
+            mVmsByName.put(name, new WeakReference<>(vm));
+            return vm;
         }
     }
 
@@ -185,14 +221,14 @@
     public VirtualMachine getOrCreate(
             @NonNull String name, @NonNull VirtualMachineConfig config)
             throws VirtualMachineException {
-        VirtualMachine vm;
-        synchronized (VirtualMachine.sCreateLock) {
-            vm = get(name);
-            if (vm == null) {
-                vm = create(name, config);
+        synchronized (sCreateLock) {
+            VirtualMachine vm = getLocked(name);
+            if (vm != null) {
+                return vm;
+            } else {
+                return createLocked(name, config);
             }
         }
-        return vm;
     }
 
     /**
@@ -208,8 +244,26 @@
      */
     public void delete(@NonNull String name) throws VirtualMachineException {
         requireNonNull(name);
-        synchronized (VirtualMachine.sCreateLock) {
-            VirtualMachine.delete(mContext, name);
+        synchronized (sCreateLock) {
+            VirtualMachine vm = getVmByName(name);
+            if (vm == null) {
+                VirtualMachine.deleteVmDirectory(mContext, name);
+            } else {
+                vm.delete(mContext, name);
+            }
+            mVmsByName.remove(name);
         }
     }
+
+    @GuardedBy("sCreateLock")
+    private VirtualMachine getVmByName(String name) {
+        WeakReference<VirtualMachine> weakReference = mVmsByName.get(name);
+        if (weakReference != null) {
+            VirtualMachine vm = weakReference.get();
+            if (vm != null && vm.getStatus() != VirtualMachine.STATUS_DELETED) {
+                return vm;
+            }
+        }
+        return null;
+    }
 }