Share a single VirtualMachine object for VMs

VirtualMachineManager.get() has always been returning a new
VirtualMachine object. This can cause confusion because VirtualMachine
object stores the VM's binder object which isn't shared among instances.
So it may be possible to run multiple guest VMs on a single instance,
which isn't advised.

This fixes it by maintaining a map from the VM to the corresponding
unique VirtualMachine object. VirtualMachineManager.get() will return an
existing object, rather than creating a new one. The mapping is removed
when the VM is deleted.

Bug: 238692795
Test: atest MicrodroidTests
Change-Id: Ie273e9118f4986aebf336b415d0e87a061228e1b
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachine.java b/javalib/src/android/system/virtualmachine/VirtualMachine.java
index 828ac9f..ae84c08 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachine.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachine.java
@@ -45,11 +45,15 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.InputStreamReader;
+import java.lang.ref.WeakReference;
 import java.nio.file.FileAlreadyExistsException;
 import java.nio.file.Files;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
+import java.util.WeakHashMap;
 import java.util.concurrent.Executor;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
@@ -65,6 +69,11 @@
  * @hide
  */
 public class VirtualMachine {
+    private static final Map<Context, Map<String, WeakReference<VirtualMachine>>> sInstances =
+            new WeakHashMap<>();
+
+    private static final Object sInstancesLock = new Object();
+
     /** Name of the directory under the files directory where all VMs created for the app exist. */
     private static final String VM_DIR = "vm";
 
@@ -159,6 +168,8 @@
 
     private final ExecutorService mExecutorService = Executors.newCachedThreadPool();
 
+    @NonNull private final Context mContext;
+
     static {
         System.loadLibrary("virtualmachine_jni");
     }
@@ -166,6 +177,7 @@
     private VirtualMachine(
             @NonNull Context context, @NonNull String name, @NonNull VirtualMachineConfig config)
             throws VirtualMachineException {
+        mContext = context;
         mPackageName = context.getPackageName();
         mName = name;
         mConfig = config;
@@ -231,6 +243,18 @@
             throw new VirtualMachineException("failed to create instance partition", e);
         }
 
+        synchronized (sInstancesLock) {
+            Map<String, WeakReference<VirtualMachine>> instancesMap;
+            if (sInstances.containsKey(context)) {
+                instancesMap = sInstances.get(context);
+            } else {
+                instancesMap = new HashMap<>();
+                sInstances.put(context, instancesMap);
+            }
+
+            instancesMap.put(name, new WeakReference<>(vm));
+        }
+
         return vm;
     }
 
@@ -249,7 +273,23 @@
             throw new VirtualMachineException(e);
         }
 
-        VirtualMachine vm = new VirtualMachine(context, name, config);
+        VirtualMachine vm;
+        synchronized (sInstancesLock) {
+            Map<String, WeakReference<VirtualMachine>> instancesMap;
+            if (sInstances.containsKey(context)) {
+                instancesMap = sInstances.get(context);
+            } else {
+                instancesMap = new HashMap<>();
+                sInstances.put(context, instancesMap);
+            }
+
+            if (instancesMap.containsKey(name)) {
+                vm = instancesMap.get(name).get();
+            } else {
+                vm = new VirtualMachine(context, name, config);
+                instancesMap.put(name, new WeakReference<>(vm));
+            }
+        }
 
         // If config file exists, but the instance image file doesn't, it means that the VM is
         // corrupted. That's different from the case that the VM doesn't exist. Throw an exception
@@ -544,6 +584,11 @@
         mInstanceFilePath.delete();
         mIdsigFilePath.delete();
         vmRootDir.delete();
+
+        synchronized (sInstancesLock) {
+            Map<String, WeakReference<VirtualMachine>> instancesMap = sInstances.get(mContext);
+            if (instancesMap != null) instancesMap.remove(mName);
+        }
     }
 
     /**
diff --git a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
index 99afe98..c2060cb 100644
--- a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
+++ b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
@@ -504,4 +504,21 @@
         assertThat(bootResult.deathReason).isEqualTo(
                 VirtualMachineCallback.DEATH_REASON_MICRODROID_INVALID_PAYLOAD_CONFIG);
     }
+
+    @Test
+    public void sameInstancesShareTheSameVmObject()
+            throws VirtualMachineException, InterruptedException, IOException {
+        VirtualMachineConfig.Builder builder =
+                mInner.newVmConfigBuilder("assets/vm_config.json");
+        VirtualMachineConfig normalConfig = builder.debugLevel(DebugLevel.NONE).build();
+        VirtualMachine vm = mInner.forceCreateNewVirtualMachine("test_vm", normalConfig);
+        VirtualMachine vm2 = mInner.getVirtualMachineManager().get("test_vm");
+        assertThat(vm).isEqualTo(vm2);
+
+        VirtualMachine newVm = mInner.forceCreateNewVirtualMachine("test_vm", normalConfig);
+        VirtualMachine newVm2 = mInner.getVirtualMachineManager().get("test_vm");
+        assertThat(newVm).isEqualTo(newVm2);
+
+        assertThat(vm).isNotEqualTo(newVm);
+    }
 }