diff --git a/apex/Android.bp b/apex/Android.bp
index 52f4384..e0ca9bf 100644
--- a/apex/Android.bp
+++ b/apex/Android.bp
@@ -123,6 +123,8 @@
         // deapexer
         "deapexer",
         "debugfs_static",
+        "blkid",
+        "fsck.erofs",
 
         // sign_virt_apex
         "avbtool",
diff --git a/apex/canned_fs_config b/apex/canned_fs_config
index 1cf63b6..ce942d3 100644
--- a/apex/canned_fs_config
+++ b/apex/canned_fs_config
@@ -1 +1 @@
-/bin/crosvm 0 2000 0755 capabilities=0x4000
+/bin/virtualizationservice 0 2000 0755 capabilities=0x1000000  # CAP_SYS_RESOURCE
diff --git a/apex/sign_virt_apex_test.sh b/apex/sign_virt_apex_test.sh
index 640a3d4..03a56ca 100644
--- a/apex/sign_virt_apex_test.sh
+++ b/apex/sign_virt_apex_test.sh
@@ -23,8 +23,11 @@
 # To access host tools
 PATH=$TEST_DIR:$PATH
 DEBUGFS=$TEST_DIR/debugfs_static
+BLKID=$TEST_DIR/blkid
+FSCKEROFS=$TEST_DIR/fsck.erofs
 
-deapexer --debugfs_path $DEBUGFS extract $TEST_DIR/com.android.virt.apex $TMP_ROOT
+deapexer --debugfs_path $DEBUGFS --blkid_path $BLKID --fsckerofs_path $FSCKEROFS \
+  extract $TEST_DIR/com.android.virt.apex $TMP_ROOT
 
 if [ "$(ls -A $TMP_ROOT/etc/fs/)" ]; then
   sign_virt_apex $TEST_DIR/test.com.android.virt.pem $TMP_ROOT
diff --git a/compos/compos_key_helper/Android.bp b/compos/compos_key_helper/Android.bp
index c9480fc..cffa1e3 100644
--- a/compos/compos_key_helper/Android.bp
+++ b/compos/compos_key_helper/Android.bp
@@ -24,6 +24,7 @@
     defaults: ["compos_key_defaults"],
     srcs: ["compos_key_main.cpp"],
 
+    header_libs: ["vm_payload_restricted_headers"],
     static_libs: [
         "libcompos_key",
     ],
diff --git a/compos/compos_key_helper/compos_key_main.cpp b/compos/compos_key_helper/compos_key_main.cpp
index 4fb0762..9417584 100644
--- a/compos/compos_key_helper/compos_key_main.cpp
+++ b/compos/compos_key_helper/compos_key_main.cpp
@@ -17,7 +17,7 @@
 #include <android-base/file.h>
 #include <android-base/logging.h>
 #include <unistd.h>
-#include <vm_payload.h>
+#include <vm_payload_restricted.h>
 
 #include <string_view>
 #include <vector>
diff --git a/compos/tests/AndroidTest.xml b/compos/tests/AndroidTest.xml
index 2a84291..f9e6837 100644
--- a/compos/tests/AndroidTest.xml
+++ b/compos/tests/AndroidTest.xml
@@ -14,6 +14,8 @@
      limitations under the License.
 -->
 <configuration description="Tests for CompOS">
+    <option name="config-descriptor:metadata" key="mainline-param" value="com.google.android.art.apex" />
+
     <target_preparer class="com.android.tradefed.targetprep.RootTargetPreparer">
         <option name="force-root" value="true" />
     </target_preparer>
diff --git a/demo/java/com/android/microdroid/demo/MainActivity.java b/demo/java/com/android/microdroid/demo/MainActivity.java
index b52ef40..df6f44e 100644
--- a/demo/java/com/android/microdroid/demo/MainActivity.java
+++ b/demo/java/com/android/microdroid/demo/MainActivity.java
@@ -145,6 +145,7 @@
 
     /** Models a virtual machine and outputs from it. */
     public static class VirtualMachineModel extends AndroidViewModel {
+        private static final String VM_NAME = "demo_vm";
         private VirtualMachine mVirtualMachine;
         private final MutableLiveData<String> mConsoleOutput = new MutableLiveData<>();
         private final MutableLiveData<String> mLogOutput = new MutableLiveData<>();
@@ -266,12 +267,12 @@
                 }
                 VirtualMachineConfig config = builder.build();
                 VirtualMachineManager vmm = VirtualMachineManager.getInstance(getApplication());
-                mVirtualMachine = vmm.getOrCreate("demo_vm", config);
+                mVirtualMachine = vmm.getOrCreate(VM_NAME, config);
                 try {
                     mVirtualMachine.setConfig(config);
                 } catch (VirtualMachineException e) {
-                    mVirtualMachine.delete();
-                    mVirtualMachine = vmm.create("demo_vm", config);
+                    vmm.delete(VM_NAME);
+                    mVirtualMachine = vmm.create(VM_NAME, config);
                 }
                 mVirtualMachine.run();
                 mVirtualMachine.setCallback(Executors.newSingleThreadExecutor(), callback);
diff --git a/javalib/Android.bp b/javalib/Android.bp
index 51dd381..cb03fa1 100644
--- a/javalib/Android.bp
+++ b/javalib/Android.bp
@@ -31,6 +31,13 @@
         // android.sysprop.*, renamed by jarjar
         "com.android.system.virtualmachine.sysprop",
     ],
+    errorprone: {
+        // We use @GuardedBy and we want a test failure if our locking isn't consistent with it.
+        enabled: true,
+        javacflags: [
+            "-Xep:GuardedBy:ERROR",
+        ],
+    },
 }
 
 prebuilt_apis {
diff --git a/javalib/src/android/system/virtualmachine/ParcelVirtualMachine.java b/javalib/src/android/system/virtualmachine/ParcelVirtualMachine.java
new file mode 100644
index 0000000..808f30a
--- /dev/null
+++ b/javalib/src/android/system/virtualmachine/ParcelVirtualMachine.java
@@ -0,0 +1,92 @@
+/*
+ * Copyright 2022 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 android.system.virtualmachine;
+
+import static java.util.Objects.requireNonNull;
+
+import android.annotation.NonNull;
+import android.os.Parcel;
+import android.os.ParcelFileDescriptor;
+import android.os.Parcelable;
+
+import com.android.internal.annotations.VisibleForTesting;
+
+/**
+ * A parcelable that captures the state of a Virtual Machine.
+ *
+ * <p>You can capture the current state of VM by creating an instance of this class with {@link
+ * VirtualMachine#toParcelVirtualMachine()}, optionally pass it to another App, and then build an
+ * identical VM with the parcel received.
+ *
+ * @hide
+ */
+public final class ParcelVirtualMachine implements Parcelable {
+    private final @NonNull ParcelFileDescriptor mConfigFd;
+    private final @NonNull ParcelFileDescriptor mInstanceImgFd;
+    // TODO(b/243129654): Add trusted storage fd once it is available.
+
+    @Override
+    public int describeContents() {
+        return CONTENTS_FILE_DESCRIPTOR;
+    }
+
+    @Override
+    public void writeToParcel(Parcel out, int flags) {
+        mConfigFd.writeToParcel(out, flags);
+        mInstanceImgFd.writeToParcel(out, flags);
+    }
+
+    public static final Parcelable.Creator<ParcelVirtualMachine> CREATOR =
+            new Parcelable.Creator<ParcelVirtualMachine>() {
+                public ParcelVirtualMachine createFromParcel(Parcel in) {
+                    return new ParcelVirtualMachine(in);
+                }
+
+                public ParcelVirtualMachine[] newArray(int size) {
+                    return new ParcelVirtualMachine[size];
+                }
+            };
+
+    /**
+     * @return File descriptor of the VM configuration file config.xml.
+     * @hide
+     */
+    @VisibleForTesting
+    public @NonNull ParcelFileDescriptor getConfigFd() {
+        return mConfigFd;
+    }
+
+    /**
+     * @return File descriptor of the instance.img of the VM.
+     * @hide
+     */
+    @VisibleForTesting
+    public @NonNull ParcelFileDescriptor getInstanceImgFd() {
+        return mInstanceImgFd;
+    }
+
+    ParcelVirtualMachine(
+            @NonNull ParcelFileDescriptor configFd, @NonNull ParcelFileDescriptor instanceImgFd) {
+        mConfigFd = configFd;
+        mInstanceImgFd = instanceImgFd;
+    }
+
+    private ParcelVirtualMachine(Parcel in) {
+        mConfigFd = requireNonNull(in.readFileDescriptor());
+        mInstanceImgFd = requireNonNull(in.readFileDescriptor());
+    }
+}
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachine.java b/javalib/src/android/system/virtualmachine/VirtualMachine.java
index e2fc33e..214e7e6 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachine.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachine.java
@@ -77,7 +77,11 @@
 import java.lang.annotation.RetentionPolicy;
 import java.lang.ref.WeakReference;
 import java.nio.file.FileAlreadyExistsException;
+import java.nio.file.FileVisitResult;
 import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.SimpleFileVisitor;
+import java.nio.file.attribute.BasicFileAttributes;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
@@ -90,17 +94,22 @@
 import java.util.zip.ZipFile;
 
 /**
- * A handle to the virtual machine. The virtual machine is local to the app which created the
- * virtual machine.
+ * Represents an VM instance, with its own configuration and state. Instances are persistent and are
+ * created or retrieved via {@link VirtualMachineManager}.
+ * <p>
+ * The {@link #run} method actually starts up the VM and allows the payload code to execute. It
+ * will continue until it exits or {@link #stop} is called. Updates on the state of the VM can
+ * be received using {@link #setCallback}. The app can communicate with the VM using
+ * {@link #connectToVsockServer} or {@link #connectVsock}.
  *
  * @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<>();
 
-    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";
 
@@ -129,7 +138,6 @@
     public static final String USE_CUSTOM_VIRTUAL_MACHINE_PERMISSION =
             "android.permission.USE_CUSTOM_VIRTUAL_MACHINE";
 
-
     /**
      * Status of a virtual machine
      *
@@ -156,9 +164,6 @@
      */
     public static final int STATUS_DELETED = 2;
 
-    /** Lock for internal synchronization. */
-    private final Object mLock = new Object();
-
     /** The package which owns this VM. */
     @NonNull private final String mPackageName;
 
@@ -166,6 +171,11 @@
     @NonNull private final String mName;
 
     /**
+     * Path to the directory containing all the files related to this VM.
+     */
+    @NonNull private final File mVmRootPath;
+
+    /**
      * Path to the config file for this VM. The config file is where the configuration is persisted.
      */
     @NonNull private final File mConfigFilePath;
@@ -187,38 +197,69 @@
     }
 
     /**
-     * List of extra apks. Apks are specified by the vm config, and corresponding idsigs are to be
-     * generated.
+     * Unmodifiable list of extra apks. Apks are specified by the vm config, and corresponding
+     * idsigs are to be generated.
      */
     @NonNull private final List<ExtraApkSpec> mExtraApks;
 
     /** Size of the instance image. 10 MB. */
     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.
+    // 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();
+
+    /** Lock protecting callbacks. */
+    private final Object mCallbackLock = new Object();
+
+
     /** The configuration that is currently associated with this VM. */
-    @NonNull private VirtualMachineConfig mConfig;
+    @GuardedBy("mLock")
+    @NonNull
+    private VirtualMachineConfig mConfig;
 
     /** Handle to the "running" VM. */
-    @Nullable private IVirtualMachine mVirtualMachine;
+    @GuardedBy("mLock")
+    @Nullable
+    private IVirtualMachine mVirtualMachine;
+
+    @GuardedBy("mLock")
+    @Nullable
+    private ParcelFileDescriptor mConsoleReader;
+
+    @GuardedBy("mLock")
+    @Nullable
+    private ParcelFileDescriptor mConsoleWriter;
+
+    @GuardedBy("mLock")
+    @Nullable
+    private ParcelFileDescriptor mLogReader;
+
+    @GuardedBy("mLock")
+    @Nullable
+    private ParcelFileDescriptor mLogWriter;
 
     /** The registered callback */
-    @GuardedBy("mLock")
+    @GuardedBy("mCallbackLock")
     @Nullable
     private VirtualMachineCallback mCallback;
 
     /** The executor on which the callback will be executed */
-    @GuardedBy("mLock")
+    @GuardedBy("mCallbackLock")
     @Nullable
     private Executor mCallbackExecutor;
 
-    @Nullable private ParcelFileDescriptor mConsoleReader;
-    @Nullable private ParcelFileDescriptor mConsoleWriter;
-
-    @Nullable private ParcelFileDescriptor mLogReader;
-    @Nullable private ParcelFileDescriptor mLogWriter;
-
-    @NonNull private final Context mContext;
-
     static {
         System.loadLibrary("virtualmachine_jni");
     }
@@ -226,130 +267,164 @@
     private VirtualMachine(
             @NonNull Context context, @NonNull String name, @NonNull VirtualMachineConfig config)
             throws VirtualMachineException {
-        mContext = context;
         mPackageName = context.getPackageName();
         mName = requireNonNull(name, "Name must not be null");
         mConfig = requireNonNull(config, "Config must not be null");
-        mConfigFilePath = getConfigFilePath(context, name);
 
-        final File vmRoot = new File(context.getFilesDir(), VM_DIR);
-        final File thisVmDir = new File(vmRoot, mName);
+        File thisVmDir = getVmDir(context, mName);
+        mVmRootPath = thisVmDir;
+        mConfigFilePath = new File(thisVmDir, CONFIG_FILE);
         mInstanceFilePath = new File(thisVmDir, INSTANCE_IMAGE_FILE);
         mIdsigFilePath = new File(thisVmDir, IDSIG_FILE);
         mExtraApks = setupExtraApks(context, config, thisVmDir);
     }
 
+    @GuardedBy("sCreateLock")
+    @NonNull
+    private static Map<String, WeakReference<VirtualMachine>> getInstancesMap(Context context) {
+        return sInstances.computeIfAbsent(context, unused -> new HashMap<>());
+    }
+
+    @NonNull
+    private static File getVmDir(Context context, String name) {
+        File vmRoot = new File(context.getDataDir(), VM_DIR);
+        return new File(vmRoot, name);
+    }
+
     /**
      * Creates a virtual machine with the given name and config. Once a virtual machine is created
-     * 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()}.
+     * 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")
     @NonNull
     static VirtualMachine create(
             @NonNull Context context, @NonNull String name, @NonNull VirtualMachineConfig config)
             throws VirtualMachineException {
-        VirtualMachine vm = new VirtualMachine(context, name, config);
+        File vmDir = getVmDir(context, name);
 
         try {
-            final File thisVmDir = vm.mConfigFilePath.getParentFile();
-            Files.createDirectories(thisVmDir.getParentFile().toPath());
+            // We don't need to undo this even if VM creation fails.
+            Files.createDirectories(vmDir.getParentFile().toPath());
 
             // The checking of the existence of this directory and the creation of it is done
             // atomically. If the directory already exists (i.e. the VM with the same name was
-            // already created), FileAlreadyExistsException is thrown
-            Files.createDirectory(thisVmDir.toPath());
-
-            try (FileOutputStream output = new FileOutputStream(vm.mConfigFilePath)) {
-                vm.mConfig.serialize(output);
-            }
+            // already created), FileAlreadyExistsException is thrown.
+            Files.createDirectory(vmDir.toPath());
         } catch (FileAlreadyExistsException e) {
             throw new VirtualMachineException("virtual machine already exists", e);
         } catch (IOException e) {
-            throw new VirtualMachineException(e);
+            throw new VirtualMachineException("failed to create directory for VM", e);
         }
 
         try {
-            vm.mInstanceFilePath.createNewFile();
-        } catch (IOException e) {
-            throw new VirtualMachineException("failed to create instance image", e);
-        }
+            VirtualMachine vm = new VirtualMachine(context, name, config);
 
-        IVirtualizationService service =
-                IVirtualizationService.Stub.asInterface(
-                        ServiceManager.waitForService(SERVICE_NAME));
-
-        try {
-            service.initializeWritablePartition(
-                    ParcelFileDescriptor.open(vm.mInstanceFilePath, MODE_READ_WRITE),
-                    INSTANCE_FILE_SIZE,
-                    PartitionType.ANDROID_VM_INSTANCE);
-        } catch (FileNotFoundException e) {
-            throw new VirtualMachineException("instance image missing", e);
-        } catch (RemoteException e) {
-            throw e.rethrowAsRuntimeException();
-        } catch (ServiceSpecificException | IllegalArgumentException e) {
-            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);
+            try (FileOutputStream output = new FileOutputStream(vm.mConfigFilePath)) {
+                config.serialize(output);
+            } catch (IOException e) {
+                throw new VirtualMachineException("failed to write VM config", e);
             }
 
-            instancesMap.put(name, new WeakReference<>(vm));
-        }
+            try {
+                vm.mInstanceFilePath.createNewFile();
+            } catch (IOException e) {
+                throw new VirtualMachineException("failed to create instance image", e);
+            }
 
-        return vm;
+            IVirtualizationService service =
+                    IVirtualizationService.Stub.asInterface(
+                            ServiceManager.waitForService(SERVICE_NAME));
+
+            try {
+                service.initializeWritablePartition(
+                        ParcelFileDescriptor.open(vm.mInstanceFilePath, MODE_READ_WRITE),
+                        INSTANCE_FILE_SIZE,
+                        PartitionType.ANDROID_VM_INSTANCE);
+            } catch (FileNotFoundException e) {
+                throw new VirtualMachineException("instance image missing", e);
+            } catch (RemoteException e) {
+                throw e.rethrowAsRuntimeException();
+            } 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
+            try {
+                deleteRecursively(vmDir);
+            } catch (IOException innerException) {
+                e.addSuppressed(innerException);
+            }
+            throw e;
+        }
     }
 
     /** Loads a virtual machine that is already created before. */
+    @GuardedBy("sCreateLock")
     @Nullable
     static VirtualMachine load(
             @NonNull Context context, @NonNull String name) throws VirtualMachineException {
-        File configFilePath = getConfigFilePath(context, name);
+        File thisVmDir = getVmDir(context, name);
+        if (!thisVmDir.exists()) {
+            // The VM doesn't exist.
+            return null;
+        }
+        File configFilePath = new File(thisVmDir, CONFIG_FILE);
         VirtualMachineConfig config;
         try (FileInputStream input = new FileInputStream(configFilePath)) {
             config = VirtualMachineConfig.from(input);
-        } catch (FileNotFoundException e) {
-            // The VM doesn't exist.
-            return null;
         } catch (IOException e) {
-            throw new VirtualMachineException(e);
+            throw new VirtualMachineException("Failed to read config file", e);
         }
 
+        Map<String, WeakReference<VirtualMachine>> instancesMap = getInstancesMap(context);
+
         VirtualMachine vm = null;
-        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();
-            }
-            if (vm == null) {
-                vm = new VirtualMachine(context, name, config);
-                instancesMap.put(name, new WeakReference<>(vm));
-            }
+        if (instancesMap.containsKey(name)) {
+            vm = instancesMap.get(name).get();
+        }
+        if (vm == null) {
+            vm = new VirtualMachine(context, name, config);
         }
 
-        // 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
-        // instead of returning null.
         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;
+        }
+
+        if (vm != null) {
+            synchronized (vm.mLock) {
+                vm.checkStopped();
+            }
+        }
+
+        try {
+            deleteRecursively(getVmDir(context, name));
+        } catch (IOException e) {
+            throw new VirtualMachineException(e);
+        }
+
+        if (instancesMap != null) instancesMap.remove(name);
+    }
+
     /**
      * Returns the name of this virtual machine. The name is unique in the package and can't be
      * changed.
@@ -372,7 +447,9 @@
      */
     @NonNull
     public VirtualMachineConfig getConfig() {
-        return mConfig;
+        synchronized (mLock) {
+            return mConfig;
+        }
     }
 
     /**
@@ -382,27 +459,70 @@
      */
     @Status
     public int getStatus() {
+        IVirtualMachine virtualMachine;
+        synchronized (mLock) {
+            virtualMachine = mVirtualMachine;
+        }
+        if (virtualMachine == null) {
+            return mVmRootPath.exists() ? STATUS_STOPPED : STATUS_DELETED;
+        } else {
+            try {
+                return stateToStatus(virtualMachine.getState());
+            } catch (RemoteException e) {
+                throw e.rethrowAsRuntimeException();
+            }
+        }
+    }
+
+    private int stateToStatus(@VirtualMachineState int state) {
+        switch (state) {
+            case VirtualMachineState.STARTING:
+            case VirtualMachineState.STARTED:
+            case VirtualMachineState.READY:
+            case VirtualMachineState.FINISHED:
+                return STATUS_RUNNING;
+            case VirtualMachineState.NOT_STARTED:
+            case VirtualMachineState.DEAD:
+            default:
+                return STATUS_STOPPED;
+        }
+    }
+
+    // Throw an appropriate exception if we have a running VM, or the VM has been deleted.
+    @GuardedBy("mLock")
+    private void checkStopped() throws VirtualMachineException {
+        if (!mVmRootPath.exists()) {
+            throw new VirtualMachineException("VM has been deleted");
+        }
+        if (mVirtualMachine == null) {
+            return;
+        }
         try {
-            if (mVirtualMachine != null) {
-                switch (mVirtualMachine.getState()) {
-                    case VirtualMachineState.NOT_STARTED:
-                        return STATUS_STOPPED;
-                    case VirtualMachineState.STARTING:
-                    case VirtualMachineState.STARTED:
-                    case VirtualMachineState.READY:
-                    case VirtualMachineState.FINISHED:
-                        return STATUS_RUNNING;
-                    case VirtualMachineState.DEAD:
-                        return STATUS_STOPPED;
+            if (stateToStatus(mVirtualMachine.getState()) != STATUS_STOPPED) {
+                throw new VirtualMachineException("VM is not in stopped state");
+            }
+        } catch (RemoteException e) {
+            throw e.rethrowAsRuntimeException();
+        }
+    }
+
+    // If we have an IVirtualMachine in the running state return it, otherwise throw.
+    @GuardedBy("mLock")
+    private IVirtualMachine getRunningVm() throws VirtualMachineException {
+        try {
+            if (mVirtualMachine != null
+                    && stateToStatus(mVirtualMachine.getState()) == STATUS_RUNNING) {
+                return mVirtualMachine;
+            } else {
+                if (!mVmRootPath.exists()) {
+                    throw new VirtualMachineException("VM has been deleted");
+                } else {
+                    throw new VirtualMachineException("VM is not in running state");
                 }
             }
         } catch (RemoteException e) {
             throw e.rethrowAsRuntimeException();
         }
-        if (!mConfigFilePath.exists()) {
-            return STATUS_DELETED;
-        }
-        return STATUS_STOPPED;
     }
 
     /**
@@ -413,7 +533,7 @@
      */
     public void setCallback(@NonNull @CallbackExecutor Executor executor,
             @NonNull VirtualMachineCallback callback) {
-        synchronized (mLock) {
+        synchronized (mCallbackLock) {
             mCallback = callback;
             mCallbackExecutor = executor;
         }
@@ -425,7 +545,7 @@
      * @hide
      */
     public void clearCallback() {
-        synchronized (mLock) {
+        synchronized (mCallbackLock) {
             mCallback = null;
             mCallbackExecutor = null;
         }
@@ -435,7 +555,7 @@
     private void executeCallback(Consumer<VirtualMachineCallback> fn) {
         final VirtualMachineCallback callback;
         final Executor executor;
-        synchronized (mLock) {
+        synchronized (mCallbackLock) {
             callback = mCallback;
             executor = mCallbackExecutor;
         }
@@ -462,117 +582,121 @@
      */
     @RequiresPermission(MANAGE_VIRTUAL_MACHINE_PERMISSION)
     public void run() throws VirtualMachineException {
-        if (getStatus() != STATUS_STOPPED) {
-            throw new VirtualMachineException(this + " is not in stopped state");
-        }
+        synchronized (mLock) {
+            checkStopped();
 
-        try {
-            mIdsigFilePath.createNewFile();
-            for (ExtraApkSpec extraApk : mExtraApks) {
-                extraApk.idsig.createNewFile();
-            }
-        } catch (IOException e) {
-            // If the file already exists, exception is not thrown.
-            throw new VirtualMachineException("failed to create idsig file", e);
-        }
-
-        IVirtualizationService service =
-                IVirtualizationService.Stub.asInterface(
-                        ServiceManager.waitForService(SERVICE_NAME));
-
-        try {
-            createVmPipes();
-
-            VirtualMachineAppConfig appConfig = getConfig().toParcel();
-            appConfig.name = mName;
-
-            // Fill the idsig file by hashing the apk
-            service.createOrUpdateIdsigFile(
-                    appConfig.apk, ParcelFileDescriptor.open(mIdsigFilePath, MODE_READ_WRITE));
-
-            for (ExtraApkSpec extraApk : mExtraApks) {
-                service.createOrUpdateIdsigFile(
-                        ParcelFileDescriptor.open(extraApk.apk, MODE_READ_ONLY),
-                        ParcelFileDescriptor.open(extraApk.idsig, MODE_READ_WRITE));
-            }
-
-            // Re-open idsig file in read-only mode
-            appConfig.idsig = ParcelFileDescriptor.open(mIdsigFilePath, MODE_READ_ONLY);
-            appConfig.instanceImage = ParcelFileDescriptor.open(mInstanceFilePath, MODE_READ_WRITE);
-            List<ParcelFileDescriptor> extraIdsigs = new ArrayList<>();
-            for (ExtraApkSpec extraApk : mExtraApks) {
-                extraIdsigs.add(ParcelFileDescriptor.open(extraApk.idsig, MODE_READ_ONLY));
-            }
-            appConfig.extraIdsigs = extraIdsigs;
-
-            android.system.virtualizationservice.VirtualMachineConfig vmConfigParcel =
-                    android.system.virtualizationservice.VirtualMachineConfig.appConfig(appConfig);
-
-            // The VM should only be observed to die once
-            AtomicBoolean onDiedCalled = new AtomicBoolean(false);
-
-            IBinder.DeathRecipient deathRecipient = () -> {
-                if (onDiedCalled.compareAndSet(false, true)) {
-                    executeCallback((cb) -> cb.onStopped(VirtualMachine.this,
-                            VirtualMachineCallback.STOP_REASON_VIRTUALIZATION_SERVICE_DIED));
+            try {
+                mIdsigFilePath.createNewFile();
+                for (ExtraApkSpec extraApk : mExtraApks) {
+                    extraApk.idsig.createNewFile();
                 }
-            };
+            } catch (IOException e) {
+                // If the file already exists, exception is not thrown.
+                throw new VirtualMachineException("failed to create idsig file", e);
+            }
 
-            mVirtualMachine = service.createVm(vmConfigParcel, mConsoleWriter, mLogWriter);
-            mVirtualMachine.registerCallback(
-                    new IVirtualMachineCallback.Stub() {
-                        @Override
-                        public void onPayloadStarted(int cid, ParcelFileDescriptor stream) {
-                            executeCallback(
-                                    (cb) -> cb.onPayloadStarted(VirtualMachine.this, stream));
-                        }
+            IVirtualizationService service =
+                    IVirtualizationService.Stub.asInterface(
+                            ServiceManager.waitForService(SERVICE_NAME));
 
-                        @Override
-                        public void onPayloadReady(int cid) {
-                            executeCallback((cb) -> cb.onPayloadReady(VirtualMachine.this));
-                        }
+            try {
+                createVmPipes();
 
-                        @Override
-                        public void onPayloadFinished(int cid, int exitCode) {
-                            executeCallback(
-                                    (cb) -> cb.onPayloadFinished(VirtualMachine.this, exitCode));
-                        }
+                VirtualMachineAppConfig appConfig = getConfig().toVsConfig();
+                appConfig.name = mName;
 
-                        @Override
-                        public void onError(int cid, int errorCode, String message) {
-                            int translatedError = getTranslatedError(errorCode);
-                            executeCallback(
-                                    (cb) -> cb.onError(VirtualMachine.this, translatedError,
-                                            message));
-                        }
+                // Fill the idsig file by hashing the apk
+                service.createOrUpdateIdsigFile(
+                        appConfig.apk, ParcelFileDescriptor.open(mIdsigFilePath, MODE_READ_WRITE));
 
-                        @Override
-                        public void onDied(int cid, int reason) {
-                            service.asBinder().unlinkToDeath(deathRecipient, 0);
-                            int translatedReason = getTranslatedReason(reason);
-                            if (onDiedCalled.compareAndSet(false, true)) {
+                for (ExtraApkSpec extraApk : mExtraApks) {
+                    service.createOrUpdateIdsigFile(
+                            ParcelFileDescriptor.open(extraApk.apk, MODE_READ_ONLY),
+                            ParcelFileDescriptor.open(extraApk.idsig, MODE_READ_WRITE));
+                }
+
+                // Re-open idsig file in read-only mode
+                appConfig.idsig = ParcelFileDescriptor.open(mIdsigFilePath, MODE_READ_ONLY);
+                appConfig.instanceImage = ParcelFileDescriptor.open(mInstanceFilePath,
+                        MODE_READ_WRITE);
+                List<ParcelFileDescriptor> extraIdsigs = new ArrayList<>();
+                for (ExtraApkSpec extraApk : mExtraApks) {
+                    extraIdsigs.add(ParcelFileDescriptor.open(extraApk.idsig, MODE_READ_ONLY));
+                }
+                appConfig.extraIdsigs = extraIdsigs;
+
+                android.system.virtualizationservice.VirtualMachineConfig vmConfigParcel =
+                        android.system.virtualizationservice.VirtualMachineConfig.appConfig(
+                                appConfig);
+
+                // The VM should only be observed to die once
+                AtomicBoolean onDiedCalled = new AtomicBoolean(false);
+
+                IBinder.DeathRecipient deathRecipient = () -> {
+                    if (onDiedCalled.compareAndSet(false, true)) {
+                        executeCallback((cb) -> cb.onStopped(VirtualMachine.this,
+                                VirtualMachineCallback.STOP_REASON_VIRTUALIZATION_SERVICE_DIED));
+                    }
+                };
+
+                mVirtualMachine = service.createVm(vmConfigParcel, mConsoleWriter, mLogWriter);
+                mVirtualMachine.registerCallback(
+                        new IVirtualMachineCallback.Stub() {
+                            @Override
+                            public void onPayloadStarted(int cid, ParcelFileDescriptor stream) {
                                 executeCallback(
-                                        (cb) -> cb.onStopped(VirtualMachine.this,
-                                                translatedReason));
+                                        (cb) -> cb.onPayloadStarted(VirtualMachine.this, stream));
+                            }
+
+                            @Override
+                            public void onPayloadReady(int cid) {
+                                executeCallback((cb) -> cb.onPayloadReady(VirtualMachine.this));
+                            }
+
+                            @Override
+                            public void onPayloadFinished(int cid, int exitCode) {
+                                executeCallback(
+                                        (cb) -> cb.onPayloadFinished(VirtualMachine.this,
+                                                exitCode));
+                            }
+
+                            @Override
+                            public void onError(int cid, int errorCode, String message) {
+                                int translatedError = getTranslatedError(errorCode);
+                                executeCallback(
+                                        (cb) -> cb.onError(VirtualMachine.this, translatedError,
+                                                message));
+                            }
+
+                            @Override
+                            public void onDied(int cid, int reason) {
+                                service.asBinder().unlinkToDeath(deathRecipient, 0);
+                                int translatedReason = getTranslatedReason(reason);
+                                if (onDiedCalled.compareAndSet(false, true)) {
+                                    executeCallback(
+                                            (cb) -> cb.onStopped(VirtualMachine.this,
+                                                    translatedReason));
+                                }
+                            }
+
+                            @Override
+                            public void onRamdump(int cid, ParcelFileDescriptor ramdump) {
+                                executeCallback(
+                                        (cb) -> cb.onRamdump(VirtualMachine.this, ramdump));
                             }
                         }
-
-                        @Override
-                        public void onRamdump(int cid, ParcelFileDescriptor ramdump) {
-                            executeCallback(
-                                    (cb) -> cb.onRamdump(VirtualMachine.this, ramdump));
-                        }
-                    }
-            );
-            service.asBinder().linkToDeath(deathRecipient, 0);
-            mVirtualMachine.start();
-        } catch (IOException | IllegalStateException | ServiceSpecificException e) {
-            throw new VirtualMachineException(e);
-        } catch (RemoteException e) {
-            throw e.rethrowAsRuntimeException();
+                );
+                service.asBinder().linkToDeath(deathRecipient, 0);
+                mVirtualMachine.start();
+            } catch (IOException | IllegalStateException | ServiceSpecificException e) {
+                throw new VirtualMachineException(e);
+            } catch (RemoteException e) {
+                throw e.rethrowAsRuntimeException();
+            }
         }
     }
 
+    @GuardedBy("mLock")
     private void createVmPipes() throws VirtualMachineException {
         try {
             if (mConsoleReader == null || mConsoleWriter == null) {
@@ -587,6 +711,207 @@
                 mLogWriter = pipe[1];
             }
         } catch (IOException e) {
+            throw new VirtualMachineException("Failed to create stream for VM", e);
+        }
+    }
+
+    /**
+     * Returns the stream object representing the console output from the virtual machine.
+     *
+     * @throws VirtualMachineException if the stream could not be created.
+     * @hide
+     */
+    @NonNull
+    public InputStream getConsoleOutputStream() throws VirtualMachineException {
+        synchronized (mLock) {
+            createVmPipes();
+            return new FileInputStream(mConsoleReader.getFileDescriptor());
+        }
+    }
+
+    /**
+     * Returns the stream object representing the log output from the virtual machine.
+     *
+     * @throws VirtualMachineException if the stream could not be created.
+     * @hide
+     */
+    @NonNull
+    public InputStream getLogOutputStream() throws VirtualMachineException {
+        synchronized (mLock) {
+            createVmPipes();
+            return new FileInputStream(mLogReader.getFileDescriptor());
+        }
+    }
+
+    /**
+     * Stops this virtual machine. Stopping a virtual machine is like pulling the plug on a real
+     * computer; the machine halts immediately. Software running on the virtual machine is not
+     * notified of the event. A stopped virtual machine can be re-started by calling {@link
+     * #run()}.
+     *
+     * @throws VirtualMachineException if the virtual machine could not be stopped.
+     * @hide
+     */
+    public void stop() throws VirtualMachineException {
+        synchronized (mLock) {
+            if (mVirtualMachine == null) {
+                throw new VirtualMachineException("VM is not running");
+            }
+            try {
+                mVirtualMachine.stop();
+                mVirtualMachine = null;
+            } catch (RemoteException e) {
+                throw e.rethrowAsRuntimeException();
+            } catch (ServiceSpecificException e) {
+                throw new VirtualMachineException(e);
+            }
+        }
+    }
+
+    /**
+     * Stops this virtual machine. See {@link #stop()}.
+     *
+     * @throws VirtualMachineException if the virtual machine could not be stopped.
+     * @hide
+     */
+    @Override
+    public void close() throws VirtualMachineException {
+        stop();
+    }
+
+    private static void deleteRecursively(File dir) throws IOException {
+        // Note: This doesn't follow symlinks, which is important. Instead they are just deleted
+        // (and Files.delete deletes the link not the target).
+        Files.walkFileTree(dir.toPath(), new SimpleFileVisitor<>() {
+            @Override
+            public FileVisitResult visitFile(Path file, BasicFileAttributes attrs)
+                    throws IOException {
+                Files.delete(file);
+                return FileVisitResult.CONTINUE;
+            }
+
+            @Override
+            public FileVisitResult postVisitDirectory(Path dir, IOException e) throws IOException {
+                // Directory is deleted after we've visited (deleted) all its contents, so it
+                // should be empty by now.
+                Files.delete(dir);
+                return FileVisitResult.CONTINUE;
+            }
+        });
+    }
+
+    /**
+     * Returns the CID of this virtual machine, if it is running.
+     *
+     * @throws VirtualMachineException if the virtual machine is not running.
+     * @hide
+     */
+    public int getCid() throws VirtualMachineException {
+        synchronized (mLock) {
+            try {
+                return getRunningVm().getCid();
+            } catch (RemoteException e) {
+                throw e.rethrowAsRuntimeException();
+            }
+        }
+    }
+
+    /**
+     * Changes the config of this virtual machine to a new one. This can be used to adjust things
+     * like the number of CPU and size of the RAM, depending on the situation (e.g. the size of the
+     * application to run on the virtual machine, etc.)
+     *
+     * The new config must be {@link VirtualMachineConfig#isCompatibleWith compatible with} the
+     * existing config.
+     *
+     * @return the old config
+     * @throws VirtualMachineException if the virtual machine is not stopped, or the new config is
+     *         incompatible.
+     * @hide
+     */
+    @NonNull
+    public VirtualMachineConfig setConfig(@NonNull VirtualMachineConfig newConfig)
+            throws VirtualMachineException {
+        synchronized (mLock) {
+            VirtualMachineConfig oldConfig = mConfig;
+            if (!oldConfig.isCompatibleWith(newConfig)) {
+                throw new VirtualMachineException("incompatible config");
+            }
+            checkStopped();
+
+            try {
+                FileOutputStream output = new FileOutputStream(mConfigFilePath);
+                newConfig.serialize(output);
+                output.close();
+            } catch (IOException e) {
+                throw new VirtualMachineException("Failed to persist config", e);
+            }
+            mConfig = newConfig;
+            return oldConfig;
+        }
+    }
+
+    @Nullable
+    private static native IBinder nativeConnectToVsockServer(IBinder vmBinder, int port);
+
+    /**
+     * Connect to a VM's binder service via vsock and return the root IBinder object. Guest VMs are
+     * expected to set up vsock servers in their payload. After the host app receives the {@link
+     * VirtualMachineCallback#onPayloadReady(VirtualMachine)}, it can use this method to
+     * establish a connection to the guest VM.
+     *
+     * @throws VirtualMachineException if the virtual machine is not running or the connection
+     *         failed.
+     * @hide
+     */
+    @NonNull
+    public IBinder connectToVsockServer(int port) throws VirtualMachineException {
+        synchronized (mLock) {
+            IBinder iBinder = nativeConnectToVsockServer(getRunningVm().asBinder(), port);
+            if (iBinder == null) {
+                throw new VirtualMachineException("Failed to connect to vsock server");
+            }
+            return iBinder;
+        }
+    }
+
+    /**
+     * Opens a vsock connection to the VM on the given port.
+     *
+     * @throws VirtualMachineException if connecting fails.
+     * @hide
+     */
+    @NonNull
+    public ParcelFileDescriptor connectVsock(int port) throws VirtualMachineException {
+        synchronized (mLock) {
+            try {
+                return getRunningVm().connectVsock(port);
+            } catch (RemoteException e) {
+                throw e.rethrowAsRuntimeException();
+            } catch (ServiceSpecificException e) {
+                throw new VirtualMachineException(e);
+            }
+        }
+    }
+
+    /**
+     * Captures the current state of the VM in a {@link ParcelVirtualMachine} instance.
+     * The VM needs to be stopped to avoid inconsistency in its state representation.
+     *
+     * @return a {@link ParcelVirtualMachine} instance that represents the VM's state.
+     * @throws VirtualMachineException if the virtual machine is not stopped, or the state could not
+     *     be captured.
+     */
+    @NonNull
+    public ParcelVirtualMachine toParcelVirtualMachine() throws VirtualMachineException {
+        synchronized (mLock) {
+            checkStopped();
+        }
+        try {
+            return new ParcelVirtualMachine(
+                ParcelFileDescriptor.open(mConfigFilePath, MODE_READ_ONLY),
+                ParcelFileDescriptor.open(mInstanceFilePath, MODE_READ_ONLY));
+        } catch (IOException e) {
             throw new VirtualMachineException(e);
         }
     }
@@ -645,187 +970,6 @@
         }
     }
 
-    /**
-     * Returns the stream object representing the console output from the virtual machine.
-     *
-     * @throws VirtualMachineException if the stream could not be created.
-     * @hide
-     */
-    @NonNull
-    public InputStream getConsoleOutputStream() throws VirtualMachineException {
-        createVmPipes();
-        return new FileInputStream(mConsoleReader.getFileDescriptor());
-    }
-
-    /**
-     * Returns the stream object representing the log output from the virtual machine.
-     *
-     * @throws VirtualMachineException if the stream could not be created.
-     * @hide
-     */
-    @NonNull
-    public InputStream getLogOutputStream() throws VirtualMachineException {
-        createVmPipes();
-        return new FileInputStream(mLogReader.getFileDescriptor());
-    }
-
-    /**
-     * Stops this virtual machine. Stopping a virtual machine is like pulling the plug on a real
-     * computer; the machine halts immediately. Software running on the virtual machine is not
-     * notified of the event. A stopped virtual machine can be re-started by calling {@link
-     * #run()}.
-     *
-     * @throws VirtualMachineException if the virtual machine could not be stopped.
-     * @hide
-     */
-    public void stop() throws VirtualMachineException {
-        if (mVirtualMachine == null) return;
-        try {
-            mVirtualMachine.stop();
-            mVirtualMachine = null;
-        } catch (RemoteException e) {
-            throw e.rethrowAsRuntimeException();
-        } catch (ServiceSpecificException e) {
-            throw new VirtualMachineException(e);
-        }
-    }
-
-    /**
-     * Stops this virtual machine. See {@link #stop()}.
-     *
-     * @throws VirtualMachineException if the virtual machine could not be stopped.
-     * @hide
-     */
-    @Override
-    public void close() throws VirtualMachineException {
-        stop();
-    }
-
-    /**
-     * Deletes this virtual machine. Deleting a virtual machine means deleting any persisted data
-     * associated with it including the per-VM secret. This is an irreversible action. A virtual
-     * machine once deleted can never be restored. A new virtual machine created with the same name
-     * and the same config is different from an already deleted virtual machine.
-     *
-     * @throws VirtualMachineException if the virtual machine is not stopped.
-     * @hide
-     */
-    public void delete() throws VirtualMachineException {
-        if (getStatus() != STATUS_STOPPED) {
-            throw new VirtualMachineException("Virtual machine is not stopped");
-        }
-        final File vmRootDir = mConfigFilePath.getParentFile();
-        for (ExtraApkSpec extraApks : mExtraApks) {
-            extraApks.idsig.delete();
-        }
-        mConfigFilePath.delete();
-        mInstanceFilePath.delete();
-        mIdsigFilePath.delete();
-        vmRootDir.delete();
-
-        synchronized (sInstancesLock) {
-            Map<String, WeakReference<VirtualMachine>> instancesMap = sInstances.get(mContext);
-            if (instancesMap != null) instancesMap.remove(mName);
-        }
-    }
-
-    /**
-     * Returns the CID of this virtual machine, if it is running.
-     *
-     * @throws VirtualMachineException if the virtual machine is not running.
-     * @hide
-     */
-    @NonNull
-    public int getCid() throws VirtualMachineException {
-        if (getStatus() != STATUS_RUNNING) {
-            throw new VirtualMachineException("VM is not running");
-        }
-        try {
-            return mVirtualMachine.getCid();
-        } catch (RemoteException e) {
-            throw e.rethrowAsRuntimeException();
-        }
-    }
-
-    /**
-     * Changes the config of this virtual machine to a new one. This can be used to adjust things
-     * like the number of CPU and size of the RAM, depending on the situation (e.g. the size of the
-     * application to run on the virtual machine, etc.)
-     *
-     * The new config must be {@link VirtualMachineConfig#isCompatibleWith compatible with} the
-     * existing config.
-     *
-     * @return the old config
-     * @throws VirtualMachineException if the virtual machine is not stopped, or the new config is
-     *         incompatible.
-     * @hide
-     */
-    @NonNull
-    public VirtualMachineConfig setConfig(@NonNull VirtualMachineConfig newConfig)
-            throws VirtualMachineException {
-        final VirtualMachineConfig oldConfig = getConfig();
-        if (!oldConfig.isCompatibleWith(newConfig)) {
-            throw new VirtualMachineException("incompatible config");
-        }
-        if (getStatus() != STATUS_STOPPED) {
-            throw new VirtualMachineException(
-                    "can't change config while virtual machine is not stopped");
-        }
-
-        try {
-            FileOutputStream output = new FileOutputStream(mConfigFilePath);
-            newConfig.serialize(output);
-            output.close();
-        } catch (IOException e) {
-            throw new VirtualMachineException(e);
-        }
-        mConfig = newConfig;
-
-        return oldConfig;
-    }
-
-    @Nullable
-    private static native IBinder nativeConnectToVsockServer(IBinder vmBinder, int port);
-
-    /**
-     * Connect to a VM's binder service via vsock and return the root IBinder object. Guest VMs are
-     * expected to set up vsock servers in their payload. After the host app receives the {@link
-     * VirtualMachineCallback#onPayloadReady(VirtualMachine)}, it can use this method to
-     * establish a connection to the guest VM.
-     *
-     * @throws VirtualMachineException if the virtual machine is not running or the connection
-     *         failed.
-     * @hide
-     */
-    @NonNull
-    public IBinder connectToVsockServer(int port) throws VirtualMachineException {
-        if (getStatus() != STATUS_RUNNING) {
-            throw new VirtualMachineException("VM is not running");
-        }
-        IBinder iBinder = nativeConnectToVsockServer(mVirtualMachine.asBinder(), port);
-        if (iBinder == null) {
-            throw new VirtualMachineException("Failed to connect to vsock server");
-        }
-        return iBinder;
-    }
-
-    /**
-     * Opens a vsock connection to the VM on the given port.
-     *
-     * @throws VirtualMachineException if connecting fails.
-     * @hide
-     */
-    @NonNull
-    public ParcelFileDescriptor connectVsock(int port) throws VirtualMachineException {
-        try {
-            return mVirtualMachine.connectVsock(port);
-        } catch (RemoteException e) {
-            throw e.rethrowAsRuntimeException();
-        } catch (ServiceSpecificException e) {
-            throw new VirtualMachineException(e);
-        }
-    }
-
     @Override
     public String toString() {
         VirtualMachineConfig config = getConfig();
@@ -916,15 +1060,9 @@
                                 new File(vmDir, EXTRA_IDSIG_FILE_PREFIX + i)));
             }
 
-            return extraApks;
+            return Collections.unmodifiableList(extraApks);
         } catch (IOException e) {
             throw new VirtualMachineException("Couldn't parse extra apks from the vm config", e);
         }
     }
-
-    private static File getConfigFilePath(@NonNull Context context, @NonNull String name) {
-        final File vmRoot = new File(context.getFilesDir(), VM_DIR);
-        final File thisVmDir = new File(vmRoot, name);
-        return new File(thisVmDir, CONFIG_FILE);
-    }
 }
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java b/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
index 2dff9bb..90b09c8 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
@@ -21,6 +21,7 @@
 import static java.util.Objects.requireNonNull;
 
 import android.annotation.IntDef;
+import android.annotation.IntRange;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.annotation.RequiresPermission;
@@ -39,15 +40,16 @@
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.util.Objects;
+
 /**
  * Represents a configuration of a virtual machine. A configuration consists of hardware
  * configurations like the number of CPUs and the size of RAM, and software configurations like the
- * OS and application to run on the virtual machine.
+ * payload to run on the virtual machine.
  *
  * @hide
  */
 public final class VirtualMachineConfig {
-    // These defines the schema of the config file persisted on disk.
+    // These define the schema of the config file persisted on disk.
     private static final int VERSION = 2;
     private static final String KEY_VERSION = "version";
     private static final String KEY_APKPATH = "apkPath";
@@ -133,6 +135,43 @@
         if (!apkPath.startsWith("/")) {
             throw new IllegalArgumentException("APK path must be an absolute path");
         }
+
+        if (memoryMib < 0) {
+            throw new IllegalArgumentException("Memory size cannot be negative");
+        }
+
+        int availableCpus = Runtime.getRuntime().availableProcessors();
+        if (numCpus < 1 || numCpus > availableCpus) {
+            throw new IllegalArgumentException("Number of vCPUs (" + numCpus + ") is out of "
+                    + "range [1, " + availableCpus + "]");
+        }
+
+        if (debugLevel != DEBUG_LEVEL_NONE && debugLevel != DEBUG_LEVEL_APP_ONLY
+                && debugLevel != DEBUG_LEVEL_FULL) {
+            throw new IllegalArgumentException("Invalid debugLevel: " + debugLevel);
+        }
+
+        if (payloadBinaryPath == null) {
+            if (payloadConfigPath == null) {
+                throw new IllegalStateException("setPayloadBinaryPath must be called");
+            }
+        } else {
+            if (payloadConfigPath != null) {
+                throw new IllegalStateException(
+                        "setPayloadBinaryPath and setPayloadConfigPath may not both be called");
+            }
+        }
+
+        if (protectedVm
+                && !HypervisorProperties.hypervisor_protected_vm_supported().orElse(false)) {
+            throw new UnsupportedOperationException(
+                    "Protected VMs are not supported on this device.");
+        }
+        if (!protectedVm && !HypervisorProperties.hypervisor_vm_supported().orElse(false)) {
+            throw new UnsupportedOperationException(
+                    "Unprotected VMs are not supported on this device.");
+        }
+
         mApkPath = apkPath;
         mPayloadConfigPath = payloadConfigPath;
         mPayloadBinaryPath = payloadBinaryPath;
@@ -177,7 +216,7 @@
     }
 
     /** Persists this config to a stream, for example a file. */
-    /* package */ void serialize(@NonNull OutputStream output) throws IOException {
+    void serialize(@NonNull OutputStream output) throws IOException {
         PersistableBundle b = new PersistableBundle();
         b.putInt(KEY_VERSION, VERSION);
         b.putString(KEY_APKPATH, mApkPath);
@@ -204,7 +243,8 @@
     }
 
     /**
-     * Returns the path to the payload config within the owning application.
+     * Returns the path within the APK to the payload config file that defines software aspects
+     * of the VM.
      *
      * @hide
      */
@@ -214,8 +254,8 @@
     }
 
     /**
-     * Returns the path within the APK to the payload binary file that will be executed within the
-     * VM.
+     * Returns the path within the {@code lib/<ABI>} directory of the APK to the payload binary
+     * file that will be executed within the VM.
      *
      * @hide
      */
@@ -249,6 +289,7 @@
      *
      * @hide
      */
+    @IntRange(from = 0)
     public int getMemoryMib() {
         return mMemoryMib;
     }
@@ -258,6 +299,7 @@
      *
      * @hide
      */
+    @IntRange(from = 1)
     public int getNumCpus() {
         return mNumCpus;
     }
@@ -279,41 +321,42 @@
     }
 
     /**
-     * Converts this config object into a parcel. Used when creating a VM via the virtualization
-     * service. Notice that the files are not passed as paths, but as file descriptors because the
-     * service doesn't accept paths as it might not have permission to open app-owned files and that
-     * could be abused to run a VM with software that the calling application doesn't own.
+     * Converts this config object into the parcelable type used when creating a VM via the
+     * virtualization service. Notice that the files are not passed as paths, but as file
+     * descriptors because the service doesn't accept paths as it might not have permission to open
+     * app-owned files and that could be abused to run a VM with software that the calling
+     * application doesn't own.
      */
-    VirtualMachineAppConfig toParcel() throws FileNotFoundException {
-        VirtualMachineAppConfig parcel = new VirtualMachineAppConfig();
-        parcel.apk = ParcelFileDescriptor.open(new File(mApkPath), MODE_READ_ONLY);
+    VirtualMachineAppConfig toVsConfig() throws FileNotFoundException {
+        VirtualMachineAppConfig vsConfig = new VirtualMachineAppConfig();
+        vsConfig.apk = ParcelFileDescriptor.open(new File(mApkPath), MODE_READ_ONLY);
         if (mPayloadBinaryPath != null) {
             VirtualMachinePayloadConfig payloadConfig = new VirtualMachinePayloadConfig();
             payloadConfig.payloadPath = mPayloadBinaryPath;
-            parcel.payload =
+            vsConfig.payload =
                     VirtualMachineAppConfig.Payload.payloadConfig(payloadConfig);
         } else {
-            parcel.payload =
+            vsConfig.payload =
                     VirtualMachineAppConfig.Payload.configPath(mPayloadConfigPath);
         }
         switch (mDebugLevel) {
             case DEBUG_LEVEL_APP_ONLY:
-                parcel.debugLevel = VirtualMachineAppConfig.DebugLevel.APP_ONLY;
+                vsConfig.debugLevel = VirtualMachineAppConfig.DebugLevel.APP_ONLY;
                 break;
             case DEBUG_LEVEL_FULL:
-                parcel.debugLevel = VirtualMachineAppConfig.DebugLevel.FULL;
+                vsConfig.debugLevel = VirtualMachineAppConfig.DebugLevel.FULL;
                 break;
             default:
-                parcel.debugLevel = VirtualMachineAppConfig.DebugLevel.NONE;
+                vsConfig.debugLevel = VirtualMachineAppConfig.DebugLevel.NONE;
                 break;
         }
-        parcel.protectedVm = mProtectedVm;
-        parcel.memoryMib = mMemoryMib;
-        parcel.numCpus = mNumCpus;
+        vsConfig.protectedVm = mProtectedVm;
+        vsConfig.memoryMib = mMemoryMib;
+        vsConfig.numCpus = mNumCpus;
         // Don't allow apps to set task profiles ... at last for now. Also, don't forget to
         // validate the string because these are appended to the cmdline argument.
-        parcel.taskProfiles = new String[0];
-        return parcel;
+        vsConfig.taskProfiles = new String[0];
+        return vsConfig;
     }
 
     /**
@@ -333,7 +376,7 @@
         private int mNumCpus;
 
         /**
-         * Creates a builder for the given context (APK).
+         * Creates a builder for the given context.
          *
          * @hide
          */
@@ -352,37 +395,10 @@
         public VirtualMachineConfig build() {
             String apkPath = (mApkPath == null) ? mContext.getPackageCodePath() : mApkPath;
 
-            int availableCpus = Runtime.getRuntime().availableProcessors();
-            if (mNumCpus < 1 || mNumCpus > availableCpus) {
-                throw new IllegalArgumentException("Number of vCPUs (" + mNumCpus + ") is out of "
-                        + "range [1, " + availableCpus + "]");
-            }
-
-            if (mPayloadBinaryPath == null) {
-                if (mPayloadConfigPath == null) {
-                    throw new IllegalStateException("payloadBinaryPath must be set");
-                }
-            } else {
-                if (mPayloadConfigPath != null) {
-                    throw new IllegalStateException(
-                            "payloadBinaryPath and payloadConfigPath may not both be set");
-                }
-            }
-
             if (!mProtectedVmSet) {
                 throw new IllegalStateException("setProtectedVm must be called explicitly");
             }
 
-            if (mProtectedVm
-                    && !HypervisorProperties.hypervisor_protected_vm_supported().orElse(false)) {
-                throw new UnsupportedOperationException(
-                        "Protected VMs are not supported on this device.");
-            }
-            if (!mProtectedVm && !HypervisorProperties.hypervisor_vm_supported().orElse(false)) {
-                throw new UnsupportedOperationException(
-                        "Unprotected VMs are not supported on this device.");
-            }
-
             return new VirtualMachineConfig(
                     apkPath, mPayloadConfigPath, mPayloadBinaryPath, mDebugLevel, mProtectedVm,
                     mMemoryMib, mNumCpus);
@@ -402,7 +418,8 @@
 
         /**
          * Sets the path within the APK to the payload config file that defines software aspects
-         * of the VM.
+         * of the VM. The file is a JSON file; see
+         * packages/modules/Virtualization/microdroid/payload/config/src/lib.rs for the format.
          *
          * @hide
          */
@@ -426,7 +443,7 @@
         }
 
         /**
-         * Sets the debug level
+         * Sets the debug level. Defaults to {@link #DEBUG_LEVEL_NONE}.
          *
          * @hide
          */
@@ -451,24 +468,25 @@
         }
 
         /**
-         * Sets the amount of RAM to give the VM. If this is zero or negative then the default will
-         * be used.
+         * Sets the amount of RAM to give the VM, in mebibytes. If zero or not explicitly set
+         * than a default size will be used.
          *
          * @hide
          */
         @NonNull
-        public Builder setMemoryMib(int memoryMib) {
+        public Builder setMemoryMib(@IntRange(from = 0) int memoryMib) {
             mMemoryMib = memoryMib;
             return this;
         }
 
         /**
-         * Sets the number of vCPUs in the VM. Defaults to 1.
+         * Sets the number of vCPUs in the VM. Defaults to 1. Cannot be more than the number of
+         * real CPUs (as returned by {@link Runtime#availableProcessors()}).
          *
          * @hide
          */
         @NonNull
-        public Builder setNumCpus(int num) {
+        public Builder setNumCpus(@IntRange(from = 1) int num) {
             mNumCpus = num;
             return this;
         }
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachineException.java b/javalib/src/android/system/virtualmachine/VirtualMachineException.java
index 88b5ea3..828775a 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachineException.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachineException.java
@@ -24,10 +24,6 @@
  * @hide
  */
 public class VirtualMachineException extends Exception {
-    public VirtualMachineException() {
-        super();
-    }
-
     public VirtualMachineException(@Nullable String message) {
         super(message);
     }
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachineManager.java b/javalib/src/android/system/virtualmachine/VirtualMachineManager.java
index 3f904be..34b9fd9 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachineManager.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachineManager.java
@@ -26,6 +26,8 @@
 import android.content.Context;
 import android.sysprop.HypervisorProperties;
 
+import com.android.internal.annotations.GuardedBy;
+
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.lang.ref.WeakReference;
@@ -33,7 +35,14 @@
 import java.util.WeakHashMap;
 
 /**
- * Manages {@link VirtualMachine} objects created for an application.
+ * Manages {@link VirtualMachine virtual machine} instances created by an app. Each instance is
+ * created from a {@link VirtualMachineConfig configuration} that defines the shape of the VM
+ * (RAM, CPUs), the code to execute within it, etc.
+ * <p>
+ * Each virtual machine instance is named; the configuration and related state of each is
+ * persisted in the app's private data directory and an instance can be retrieved given the name.
+ * <p>
+ * The app can then start, stop and otherwise interact with the VM.
  *
  * @hide
  */
@@ -44,6 +53,7 @@
         mContext = context;
     }
 
+    @GuardedBy("sInstances")
     private static final Map<Context, WeakReference<VirtualMachineManager>> sInstances =
             new WeakHashMap<>();
 
@@ -89,9 +99,6 @@
         }
     }
 
-    /** A lock used to synchronize the creation of virtual machines */
-    private static final Object sCreateLock = new Object();
-
     /**
      * Returns a set of flags indicating what this implementation of virtualization is capable of.
      *
@@ -129,7 +136,7 @@
     public VirtualMachine create(
             @NonNull String name, @NonNull VirtualMachineConfig config)
             throws VirtualMachineException {
-        synchronized (sCreateLock) {
+        synchronized (VirtualMachine.sCreateLock) {
             return VirtualMachine.create(mContext, name, config);
         }
     }
@@ -138,12 +145,15 @@
      * Returns an existing {@link VirtualMachine} with the given name. Returns null if there is no
      * such virtual machine.
      *
-     * @throws VirtualMachineException if the virtual machine could not be successfully retrieved.
+     * @throws VirtualMachineException if the virtual machine exists but could not be successfully
+     *                                 retrieved.
      * @hide
      */
     @Nullable
     public VirtualMachine get(@NonNull String name) throws VirtualMachineException {
-        return VirtualMachine.load(mContext, name);
+        synchronized (VirtualMachine.sCreateLock) {
+            return VirtualMachine.load(mContext, name);
+        }
     }
 
     /**
@@ -158,7 +168,7 @@
             @NonNull String name, @NonNull VirtualMachineConfig config)
             throws VirtualMachineException {
         VirtualMachine vm;
-        synchronized (sCreateLock) {
+        synchronized (VirtualMachine.sCreateLock) {
             vm = get(name);
             if (vm == null) {
                 vm = create(name, config);
@@ -166,4 +176,22 @@
         }
         return vm;
     }
+
+    /**
+     * Deletes an existing {@link VirtualMachine}. Deleting a virtual machine means deleting any
+     * persisted data associated with it including the per-VM secret. This is an irreversible
+     * action. A virtual machine once deleted can never be restored. A new virtual machine created
+     * with the same name is different from an already deleted virtual machine even if it has the
+     * same config.
+     *
+     * @throws VirtualMachineException if the virtual machine does not exist, is not stopped,
+     *                                 or cannot be deleted.
+     * @hide
+     */
+    public void delete(@NonNull String name) throws VirtualMachineException {
+        requireNonNull(name);
+        synchronized (VirtualMachine.sCreateLock) {
+            VirtualMachine.delete(mContext, name);
+        }
+    }
 }
diff --git a/microdroid/kernel/README.md b/microdroid/kernel/README.md
index 1f6b2ee..38eca71 100644
--- a/microdroid/kernel/README.md
+++ b/microdroid/kernel/README.md
@@ -24,6 +24,7 @@
 
 For x86\_64,
 ```bash
+tools/bazel clean
 tools/bazel run --config=fast --lto=thin //common-modules/virtual-device:microdroid_x86_64_dist -- --dist_dir=out/dist
 ```
 
@@ -36,7 +37,15 @@
 
 ### Change the kernel configs
 
+For ARM64
+```bash
+tools/bazel run //common-modules/virtual-device:microdroid_aarch64_config menuconfig
+```
 
+For x86\_64
+```bash
+tools/bazel run //common-modules/virtual-device:microdroid_x86_64_config menuconfig
+```
 
 ## How to update Microdroid kernel prebuilts
 
diff --git a/microdroid/vm_payload/Android.bp b/microdroid/vm_payload/Android.bp
index 8d78444..e153f92 100644
--- a/microdroid/vm_payload/Android.bp
+++ b/microdroid/vm_payload/Android.bp
@@ -29,7 +29,7 @@
 
 rust_bindgen {
     name: "libvm_payload_bindgen",
-    wrapper_src: "include/vm_payload.h",
+    wrapper_src: "include-restricted/vm_payload_restricted.h",
     crate_name: "vm_payload_bindgen",
     source_stem: "bindings",
     apex_available: ["com.android.compos"],
@@ -41,5 +41,15 @@
 
 cc_library_headers {
     name: "vm_payload_headers",
+    apex_available: ["com.android.compos"],
     export_include_dirs: ["include"],
 }
+
+cc_library_headers {
+    name: "vm_payload_restricted_headers",
+    header_libs: ["vm_payload_headers"],
+    export_header_lib_headers: ["vm_payload_headers"],
+    export_include_dirs: ["include-restricted"],
+    apex_available: ["com.android.compos"],
+    visibility: ["//packages/modules/Virtualization:__subpackages__"],
+}
diff --git a/microdroid/vm_payload/include-restricted/vm_payload_restricted.h b/microdroid/vm_payload/include-restricted/vm_payload_restricted.h
new file mode 100644
index 0000000..8170a64
--- /dev/null
+++ b/microdroid/vm_payload/include-restricted/vm_payload_restricted.h
@@ -0,0 +1,56 @@
+/*
+ * Copyright (C) 2022 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.
+ */
+
+#pragma once
+
+#include <stdbool.h>
+#include <stddef.h>
+#include <sys/cdefs.h>
+
+#include "vm_payload.h"
+
+// The functions declared here are restricted to VMs created with a config file;
+// they will fail if called in other VMs. The ability to create such VMs
+// requires the android.permission.USE_CUSTOM_VIRTUAL_MACHINE permission, and is
+// therefore not available to privileged or third party apps.
+
+// These functions can be used by tests, if the permission is granted via shell.
+
+__BEGIN_DECLS
+
+/**
+ * Get the VM's DICE attestation chain.
+ *
+ * \param data pointer to size bytes where the chain is written.
+ * \param size number of bytes that can be written to data.
+ * \param total outputs the total size of the chain if the function succeeds
+ *
+ * \return true on success and false on failure.
+ */
+bool AVmPayload_getDiceAttestationChain(void *data, size_t size, size_t *total);
+
+/**
+ * Get the VM's DICE attestation CDI.
+ *
+ * \param data pointer to size bytes where the CDI is written.
+ * \param size number of bytes that can be written to data.
+ * \param total outputs the total size of the CDI if the function succeeds
+ *
+ * \return true on success and false on failure.
+ */
+bool AVmPayload_getDiceAttestationCdi(void *data, size_t size, size_t *total);
+
+__END_DECLS
diff --git a/microdroid/vm_payload/include/vm_payload.h b/microdroid/vm_payload/include/vm_payload.h
index 2aeb44e..82dbd6d 100644
--- a/microdroid/vm_payload/include/vm_payload.h
+++ b/microdroid/vm_payload/include/vm_payload.h
@@ -18,12 +18,11 @@
 
 #include <stdbool.h>
 #include <stddef.h>
+#include <sys/cdefs.h>
 
 #include "vm_main.h"
 
-#ifdef __cplusplus
-extern "C" {
-#endif
+__BEGIN_DECLS
 
 struct AIBinder;
 typedef struct AIBinder AIBinder;
@@ -71,32 +70,6 @@
                                     size_t size);
 
 /**
- * Get the VM's DICE attestation chain.
- *
- * This function will fail if the use of restricted APIs is not permitted.
- *
- * \param data pointer to size bytes where the chain is written.
- * \param size number of bytes that can be written to data.
- * \param total outputs the total size of the chain if the function succeeds
- *
- * \return true on success and false on failure.
- */
-bool AVmPayload_getDiceAttestationChain(void *data, size_t size, size_t *total);
-
-/**
- * Get the VM's DICE attestation CDI.
- *
- * This function will fail if the use of restricted APIs is not permitted.
- *
- * \param data pointer to size bytes where the CDI is written.
- * \param size number of bytes that can be written to data.
- * \param total outputs the total size of the CDI if the function succeeds
- *
- * \return true on success and false on failure.
- */
-bool AVmPayload_getDiceAttestationCdi(void *data, size_t size, size_t *total);
-
-/**
  * Gets the path to the APK contents. It is a directory, under which are
  * the unzipped contents of the APK containing the payload, all read-only
  * but accessible to the payload.
@@ -107,6 +80,4 @@
  */
 const char *AVmPayload_getApkContentsPath(void);
 
-#ifdef __cplusplus
-} // extern "C"
-#endif
+__END_DECLS
diff --git a/microdroid/vm_payload/src/vm_payload_service.rs b/microdroid/vm_payload/src/vm_payload_service.rs
index b0dd891..098d246 100644
--- a/microdroid/vm_payload/src/vm_payload_service.rs
+++ b/microdroid/vm_payload/src/vm_payload_service.rs
@@ -15,12 +15,12 @@
 //! This module handles the interaction with virtual machine payload service.
 
 use android_system_virtualization_payload::aidl::android::system::virtualization::payload::IVmPayloadService::{
-    IVmPayloadService, VM_PAYLOAD_SERVICE_NAME, VM_APK_CONTENTS_PATH};
+    IVmPayloadService, VM_PAYLOAD_SERVICE_SOCKET_NAME, VM_APK_CONTENTS_PATH};
 use anyhow::{Context, Result};
-use binder::{wait_for_interface, Strong, unstable_api::{AIBinder, new_spibinder}};
+use binder::{Strong, unstable_api::{AIBinder, new_spibinder}};
 use lazy_static::lazy_static;
 use log::{error, info, Level};
-use rpcbinder::run_vsock_rpc_server;
+use rpcbinder::{get_unix_domain_rpc_interface, run_vsock_rpc_server};
 use std::ffi::CString;
 use std::os::raw::{c_char, c_void};
 
@@ -203,6 +203,6 @@
 }
 
 fn get_vm_payload_service() -> Result<Strong<dyn IVmPayloadService>> {
-    wait_for_interface(VM_PAYLOAD_SERVICE_NAME)
-        .context(format!("Failed to connect to service: {}", VM_PAYLOAD_SERVICE_NAME))
+    get_unix_domain_rpc_interface(VM_PAYLOAD_SERVICE_SOCKET_NAME)
+        .context(format!("Failed to connect to service: {}", VM_PAYLOAD_SERVICE_SOCKET_NAME))
 }
diff --git a/microdroid_manager/aidl/android/system/virtualization/payload/IVmPayloadService.aidl b/microdroid_manager/aidl/android/system/virtualization/payload/IVmPayloadService.aidl
index 4823bb8..f8e7d34 100644
--- a/microdroid_manager/aidl/android/system/virtualization/payload/IVmPayloadService.aidl
+++ b/microdroid_manager/aidl/android/system/virtualization/payload/IVmPayloadService.aidl
@@ -21,8 +21,8 @@
  * Microdroid Manager for execution.
  */
 interface IVmPayloadService {
-    /** Name of the service IVmPayloadService. */
-    const String VM_PAYLOAD_SERVICE_NAME = "virtual_machine_payload_service";
+    /** Socket name of the service IVmPayloadService. */
+    const String VM_PAYLOAD_SERVICE_SOCKET_NAME = "vm_payload_service";
 
     /** Path to the APK contents path. */
     const String VM_APK_CONTENTS_PATH = "/mnt/apk";
diff --git a/microdroid_manager/microdroid_manager.rc b/microdroid_manager/microdroid_manager.rc
index 74a219d..cfa70bd 100644
--- a/microdroid_manager/microdroid_manager.rc
+++ b/microdroid_manager/microdroid_manager.rc
@@ -6,3 +6,4 @@
     oneshot
     # SYS_BOOT is required to exec kexecload from microdroid_manager
     capabilities AUDIT_CONTROL SYS_ADMIN SYS_BOOT
+    socket vm_payload_service stream 0666 system system
diff --git a/microdroid_manager/src/main.rs b/microdroid_manager/src/main.rs
index b8e85e7..4b4f996 100644
--- a/microdroid_manager/src/main.rs
+++ b/microdroid_manager/src/main.rs
@@ -18,7 +18,6 @@
 mod instance;
 mod ioutil;
 mod payload;
-mod procutil;
 mod swap;
 mod vm_payload_service;
 
@@ -26,17 +25,13 @@
 use crate::instance::{ApexData, ApkData, InstanceDisk, MicrodroidData, RootHash};
 use crate::vm_payload_service::register_vm_payload_service;
 use android_system_virtualizationcommon::aidl::android::system::virtualizationcommon::ErrorCode::ErrorCode;
-use android_system_virtualmachineservice::aidl::android::system::virtualmachineservice::{
-    IVirtualMachineService::{
+use android_system_virtualmachineservice::aidl::android::system::virtualmachineservice::IVirtualMachineService::{
         IVirtualMachineService, VM_BINDER_SERVICE_PORT, VM_STREAM_SERVICE_PORT,
-    },
-    VirtualMachineCpuStatus::VirtualMachineCpuStatus,
-    VirtualMachineMemStatus::VirtualMachineMemStatus,
 };
 use android_system_virtualization_payload::aidl::android::system::virtualization::payload::IVmPayloadService::VM_APK_CONTENTS_PATH;
 use anyhow::{anyhow, bail, ensure, Context, Error, Result};
 use apkverify::{get_public_key_der, verify, V4Signature};
-use binder::{ProcessState, Strong};
+use binder::Strong;
 use diced_utils::cbor::{encode_header, encode_number};
 use glob::glob;
 use itertools::sorted;
@@ -46,7 +41,6 @@
 use microdroid_payload_config::{OsConfig, Task, TaskType, VmPayloadConfig};
 use openssl::sha::Sha512;
 use payload::{get_apex_data_from_payload, load_metadata, to_metadata};
-use procutil::{get_cpu_time, get_mem_info};
 use rand::Fill;
 use rpcbinder::get_vsock_rpc_interface;
 use rustutils::system_properties;
@@ -59,12 +53,10 @@
 use std::path::Path;
 use std::process::{Child, Command, Stdio};
 use std::str;
-use std::thread;
 use std::time::{Duration, SystemTime};
 use vsock::VsockStream;
 
 const WAIT_TIMEOUT: Duration = Duration::from_secs(10);
-const SENDING_VM_STATUS_CYCLE_PERIOD: Duration = Duration::from_secs(60);
 const MAIN_APK_PATH: &str = "/dev/block/by-name/microdroid-apk";
 const MAIN_APK_IDSIG_PATH: &str = "/dev/block/by-name/microdroid-apk-idsig";
 const MAIN_APK_DEVICE_NAME: &str = "microdroid-apk";
@@ -97,42 +89,6 @@
     InvalidConfig(String),
 }
 
-fn send_vm_status(service: &Strong<dyn IVirtualMachineService>) -> Result<()> {
-    // Collect VM CPU time information and creating VmCpuStatus atom for metrics.
-    let cpu_time = get_cpu_time()?;
-    let vm_cpu_status = VirtualMachineCpuStatus {
-        cpu_time_user: cpu_time.user,
-        cpu_time_nice: cpu_time.nice,
-        cpu_time_sys: cpu_time.sys,
-        cpu_time_idle: cpu_time.idle,
-    };
-    service.notifyCpuStatus(&vm_cpu_status).expect("Can't send information about VM CPU status");
-
-    // Collect VM memory information and creating VmMemStatus atom for metrics.
-    let mem_info = get_mem_info()?;
-    let vm_mem_status = VirtualMachineMemStatus {
-        mem_total: mem_info.total,
-        mem_free: mem_info.free,
-        mem_available: mem_info.available,
-        mem_buffer: mem_info.buffer,
-        mem_cached: mem_info.cached,
-    };
-    service.notifyMemStatus(&vm_mem_status).expect("Can't send information about VM memory status");
-
-    Ok(())
-}
-
-fn send_vm_status_periodically() -> Result<()> {
-    let service = get_vms_rpc_binder()
-        .context("cannot connect to VirtualMachineService")
-        .map_err(|e| MicrodroidError::FailedToConnectToVirtualizationService(e.to_string()))?;
-
-    loop {
-        send_vm_status(&service)?;
-        thread::sleep(SENDING_VM_STATUS_CYCLE_PERIOD);
-    }
-}
-
 fn translate_error(err: &Error) -> (ErrorCode, String) {
     if let Some(e) = err.downcast_ref::<MicrodroidError>() {
         match e {
@@ -225,12 +181,6 @@
         .context("cannot connect to VirtualMachineService")
         .map_err(|e| MicrodroidError::FailedToConnectToVirtualizationService(e.to_string()))?;
 
-    thread::spawn(move || {
-        if let Err(e) = send_vm_status_periodically() {
-            error!("failed to get virtual machine status: {:?}", e);
-        }
-    });
-
     match try_run_payload(&service) {
         Ok(code) => {
             info!("notifying payload finished");
@@ -447,11 +397,8 @@
     wait_for_property_true(APK_MOUNT_DONE_PROP).context("Failed waiting for APK mount done")?;
 
     register_vm_payload_service(allow_restricted_apis, service.clone(), dice)?;
-    ProcessState::start_thread_pool();
 
     system_properties::write("dev.bootcomplete", "1").context("set dev.bootcomplete")?;
-    send_vm_status(service)?;
-
     exec_task(task, service).context("Failed to run payload")
 }
 
@@ -783,7 +730,6 @@
     service.notifyPayloadStarted()?;
 
     let exit_status = command.spawn()?.wait()?;
-    send_vm_status(service)?;
     exit_status.code().ok_or_else(|| anyhow!("Failed to get exit_code from the paylaod."))
 }
 
diff --git a/microdroid_manager/src/procutil.rs b/microdroid_manager/src/procutil.rs
deleted file mode 100644
index f9479c4..0000000
--- a/microdroid_manager/src/procutil.rs
+++ /dev/null
@@ -1,127 +0,0 @@
-// Copyright 2022, 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.
-
-use anyhow::{bail, Result};
-use libc::{sysconf, _SC_CLK_TCK};
-use std::fs::File;
-use std::io::{BufRead, BufReader};
-
-const MILLIS_PER_SEC: i64 = 1000;
-
-pub struct CpuTime {
-    pub user: i64,
-    pub nice: i64,
-    pub sys: i64,
-    pub idle: i64,
-}
-
-pub struct MemInfo {
-    pub total: i64,
-    pub free: i64,
-    pub available: i64,
-    pub buffer: i64,
-    pub cached: i64,
-}
-
-// Get CPU time information from /proc/stat
-//
-// /proc/stat example(omitted):
-//   cpu  24790952 21104390 10771070 10480973587 1700955 0 410931 0 316532 0
-//   cpu0 169636 141307 61153 81785791 9605 0 183524 0 1345 0
-//   cpu1 182431 198327 68273 81431817 10445 0 32392 0 2616 0
-//   cpu2 183209 174917 68591 81933935 12239 0 10042 0 2415 0
-//   cpu3 183413 177758 69908 81927474 13354 0 5853 0 2491 0
-//   intr 7913477443 39 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
-//   ctxt 10326710014
-//   btime 1664123605
-//   processes 9225712
-//   procs_running 1
-//   procs_blocked 0
-//   softirq 2683914305 14595298 304837101 1581 327291100 16397051 0 208857783 1024640365 787932 786506094
-//
-// expected output:
-//   user: 24790952
-//   nice: 21104390
-//   sys: 10771070
-//   idle: 10480973587
-pub fn get_cpu_time() -> Result<CpuTime> {
-    let mut proc_stat = BufReader::new(File::open("/proc/stat")?);
-    let mut line = String::new();
-    proc_stat.read_line(&mut line)?;
-    let data_list: Vec<_> = line.split_whitespace().filter_map(|s| s.parse::<i64>().ok()).collect();
-    if data_list.len() < 4 {
-        bail!("Failed to extract numeric values in /proc/stat :\n{}", line);
-    }
-
-    let ticks_per_sec = unsafe { sysconf(_SC_CLK_TCK) } as i64;
-    let cpu_time = CpuTime {
-        user: data_list[0] * MILLIS_PER_SEC / ticks_per_sec,
-        nice: data_list[1] * MILLIS_PER_SEC / ticks_per_sec,
-        sys: data_list[2] * MILLIS_PER_SEC / ticks_per_sec,
-        idle: data_list[3] * MILLIS_PER_SEC / ticks_per_sec,
-    };
-    Ok(cpu_time)
-}
-
-// Get memory information from /proc/meminfo
-//
-// /proc/meminfo example(omitted):
-//   MemTotal:       263742736 kB
-//   MemFree:        37144204 kB
-//   MemAvailable:   249168700 kB
-//   Buffers:        10231296 kB
-//   Cached:         189502836 kB
-//   SwapCached:       113848 kB
-//   Active:         132266424 kB
-//   Inactive:       73587504 kB
-//   Active(anon):    1455240 kB
-//   Inactive(anon):  6993584 kB
-//   Active(file):   130811184 kB
-//   Inactive(file): 66593920 kB
-//   Unevictable:       56436 kB
-//   Mlocked:           56436 kB
-//   SwapTotal:      255123452 kB
-//   SwapFree:       254499068 kB
-//   Dirty:               596 kB
-//   Writeback:             0 kB
-//   AnonPages:       5295864 kB
-//   Mapped:          3512608 kB
-//
-// expected output:
-//   total: 263742736
-//   free: 37144204
-//   available: 249168700
-//   buffer: 10231296
-//   cached: 189502836
-pub fn get_mem_info() -> Result<MemInfo> {
-    let mut proc_stat = BufReader::new(File::open("/proc/meminfo")?);
-    let mut lines = String::new();
-    for _ in 0..5 {
-        proc_stat.read_line(&mut lines)?;
-    }
-    let data_list: Vec<_> =
-        lines.split_whitespace().filter_map(|s| s.parse::<i64>().ok()).collect();
-    if data_list.len() != 5 {
-        bail!("Failed to extract numeric values in /proc/meminfo :\n{}", lines);
-    }
-
-    let mem_info = MemInfo {
-        total: data_list[0],
-        free: data_list[1],
-        available: data_list[2],
-        buffer: data_list[3],
-        cached: data_list[4],
-    };
-    Ok(mem_info)
-}
diff --git a/microdroid_manager/src/swap.rs b/microdroid_manager/src/swap.rs
index 0fd1468..d7916db 100644
--- a/microdroid_manager/src/swap.rs
+++ b/microdroid_manager/src/swap.rs
@@ -43,20 +43,21 @@
     let sysfs_size = format!("/sys/{}/size", dev);
     let len = read_to_string(&sysfs_size)?
         .trim()
-        .parse::<u32>()
-        .context(format!("No u32 in {}", &sysfs_size))?
-        * 512;
+        .parse::<u64>()
+        .context(format!("No u64 in {}", &sysfs_size))?
+        .checked_mul(512)
+        .ok_or_else(|| anyhow!("sysfs_size too large"))?;
 
-    let pagesize: libc::c_uint;
     // safe because we give a constant and known-valid sysconf parameter
-    unsafe {
-        pagesize = libc::sysconf(libc::_SC_PAGE_SIZE) as libc::c_uint;
-    }
+    let pagesize = unsafe { libc::sysconf(libc::_SC_PAGE_SIZE) as u64 };
 
     let mut f = OpenOptions::new().read(false).write(true).open(format!("/dev/{}", dev))?;
 
+    let last_page = len / pagesize - 1;
+
     // Write the info fields: [ version, last_page ]
-    let info: [u32; 2] = [1, (len / pagesize) - 1];
+    let info: [u32; 2] = [1, last_page.try_into().context("Number of pages out of range")?];
+
     f.seek(SeekFrom::Start(1024))?;
     f.write_all(&info.iter().flat_map(|v| v.to_ne_bytes()).collect::<Vec<u8>>())?;
 
diff --git a/microdroid_manager/src/vm_payload_service.rs b/microdroid_manager/src/vm_payload_service.rs
index 159bf67..fcfc79d 100644
--- a/microdroid_manager/src/vm_payload_service.rs
+++ b/microdroid_manager/src/vm_payload_service.rs
@@ -16,13 +16,17 @@
 
 use crate::dice::DiceContext;
 use android_system_virtualization_payload::aidl::android::system::virtualization::payload::IVmPayloadService::{
-    BnVmPayloadService, IVmPayloadService, VM_PAYLOAD_SERVICE_NAME};
+    BnVmPayloadService, IVmPayloadService, VM_PAYLOAD_SERVICE_SOCKET_NAME};
 use android_system_virtualmachineservice::aidl::android::system::virtualmachineservice::IVirtualMachineService::IVirtualMachineService;
-use anyhow::{Context, Result};
-use binder::{Interface, BinderFeatures, ExceptionCode, Status, Strong, add_service};
-use log::error;
+use anyhow::{bail, Result};
+use binder::{Interface, BinderFeatures, ExceptionCode, Status, Strong};
+use log::{error, info};
 use openssl::hkdf::hkdf;
 use openssl::md::Md;
+use rpcbinder::run_init_unix_domain_rpc_server;
+use std::sync::mpsc;
+use std::thread;
+use std::time::Duration;
 
 /// Implementation of `IVmPayloadService`.
 struct VmPayloadService {
@@ -97,8 +101,29 @@
         VmPayloadService::new(allow_restricted_apis, vm_service, dice),
         BinderFeatures::default(),
     );
-    add_service(VM_PAYLOAD_SERVICE_NAME, vm_payload_binder.as_binder())
-        .with_context(|| format!("Failed to register service {}", VM_PAYLOAD_SERVICE_NAME))?;
-    log::info!("{} is running", VM_PAYLOAD_SERVICE_NAME);
-    Ok(())
+    let (sender, receiver) = mpsc::channel();
+    thread::spawn(move || {
+        let retval = run_init_unix_domain_rpc_server(
+            vm_payload_binder.as_binder(),
+            VM_PAYLOAD_SERVICE_SOCKET_NAME,
+            || {
+                sender.send(()).unwrap();
+            },
+        );
+        if retval {
+            info!(
+                "The RPC server at '{}' has shut down gracefully.",
+                VM_PAYLOAD_SERVICE_SOCKET_NAME
+            );
+        } else {
+            error!("Premature termination of the RPC server '{}'.", VM_PAYLOAD_SERVICE_SOCKET_NAME);
+        }
+    });
+    match receiver.recv_timeout(Duration::from_millis(200)) {
+        Ok(()) => {
+            info!("The RPC server '{}' is running.", VM_PAYLOAD_SERVICE_SOCKET_NAME);
+            Ok(())
+        }
+        _ => bail!("Failed to register service '{}'", VM_PAYLOAD_SERVICE_SOCKET_NAME),
+    }
 }
diff --git a/tests/benchmark_hostside/Android.bp b/tests/benchmark_hostside/Android.bp
index c5c7571..e053406 100644
--- a/tests/benchmark_hostside/Android.bp
+++ b/tests/benchmark_hostside/Android.bp
@@ -17,4 +17,7 @@
     test_suites: [
         "general-tests",
     ],
+    data: [
+        ":MicrodroidTestApp",
+    ],
 }
diff --git a/tests/benchmark_hostside/java/android/avf/test/AVFHostTestCase.java b/tests/benchmark_hostside/java/android/avf/test/AVFHostTestCase.java
index 5e57b32..1996f4b 100644
--- a/tests/benchmark_hostside/java/android/avf/test/AVFHostTestCase.java
+++ b/tests/benchmark_hostside/java/android/avf/test/AVFHostTestCase.java
@@ -16,6 +16,7 @@
 
 package android.avf.test;
 
+import static com.android.tradefed.device.TestDevice.MicrodroidBuilder;
 import static com.android.tradefed.testtype.DeviceJUnit4ClassRunner.TestMetrics;
 
 import static com.google.common.truth.Truth.assertWithMessage;
@@ -31,6 +32,8 @@
 import com.android.microdroid.test.host.CommandRunner;
 import com.android.microdroid.test.host.MicrodroidHostTestCaseBase;
 import com.android.tradefed.device.DeviceNotAvailableException;
+import com.android.tradefed.device.ITestDevice;
+import com.android.tradefed.device.TestDevice;
 import com.android.tradefed.log.LogUtil.CLog;
 import com.android.tradefed.testtype.DeviceJUnit4ClassRunner;
 import com.android.tradefed.util.CommandResult;
@@ -45,7 +48,6 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.Optional;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
@@ -85,6 +87,9 @@
     @Before
     public void setUp() throws Exception {
         testIfDeviceIsCapable(getDevice());
+
+        getDevice().installPackage(findTestFile(APK_NAME), /* reinstall */ false);
+
         mMetricsProcessor = new MetricsProcessor(getMetricPrefix() + "hostside/");
     }
 
@@ -196,10 +201,6 @@
         return null;
     }
 
-    private void microdroidWaitForBootComplete() {
-        runOnMicrodroidForResult("watch -e \"getprop dev.bootcomplete | grep '^0$'\"");
-    }
-
     private AmStartupTimeCmdParser getColdRunStartupTimes(CommandRunner android, String pkgName)
             throws DeviceNotAvailableException, InterruptedException {
         unlockScreen(android);
@@ -216,9 +217,7 @@
     // and the time measured after running the VM.
     private void getAppStartupTime(String pkgName, StartupTimeMetricCollection metricColector)
             throws Exception {
-        final String configPath = "assets/vm_config.json";
-        final String cid;
-        final int vm_mem_mb;
+        TestDevice device = (TestDevice) getDevice();
 
         // 1. Reboot the device to run the test without stage2 fragmentation
         getDevice().rebootUntilOnline();
@@ -240,30 +239,30 @@
         android.tryRun("rm", "-rf", MicrodroidHostTestCaseBase.TEST_ROOT);
 
         // Donate 80% of the available device memory to the VM
-        vm_mem_mb = getFreeMemoryInfoMb(android) * 80 / 100;
-        cid = startMicrodroid(
-                            getDevice(),
-                            getBuild(),
-                            APK_NAME,
-                            PACKAGE_NAME,
-                            configPath,
-                            true,
-                            vm_mem_mb,
-                            Optional.of(NUM_VCPUS));
-        adbConnectToMicrodroid(getDevice(), cid);
-        microdroidWaitForBootComplete();
+        final String configPath = "assets/vm_config.json";
+        final int vm_mem_mb = getFreeMemoryInfoMb(android) * 80 / 100;
+        ITestDevice microdroidDevice =
+                MicrodroidBuilder.fromDevicePath(getPathForPackage(PACKAGE_NAME), configPath)
+                        .debugLevel("full")
+                        .memoryMib(vm_mem_mb)
+                        .numCpus(NUM_VCPUS)
+                        .build(device);
+        microdroidDevice.waitForBootComplete(30000);
+        microdroidDevice.enableAdbRoot();
 
-        rootMicrodroid();
+        CommandRunner microdroid = new CommandRunner(microdroidDevice);
 
-        runOnMicrodroid("mkdir -p /mnt/ramdisk && chmod 777 /mnt/ramdisk");
-        runOnMicrodroid("mount -t tmpfs -o size=32G tmpfs /mnt/ramdisk");
+        microdroid.run("mkdir -p /mnt/ramdisk && chmod 777 /mnt/ramdisk");
+        microdroid.run("mount -t tmpfs -o size=32G tmpfs /mnt/ramdisk");
 
         // Allocate memory for the VM until it fails and make sure that we touch
         // the allocated memory in the guest to be able to create stage2 fragmentation.
         try {
-            runOnMicrodroidForResult(String.format("cd /mnt/ramdisk && truncate -s %dM sprayMemory"
-                    + " && dd if=/dev/zero of=sprayMemory bs=1MB count=%d",
-                    vm_mem_mb , vm_mem_mb));
+            microdroid.tryRun(
+                    String.format(
+                            "cd /mnt/ramdisk && truncate -s %dM sprayMemory"
+                                    + " && dd if=/dev/zero of=sprayMemory bs=1MB count=%d",
+                            vm_mem_mb, vm_mem_mb));
         } catch (Exception ex) {
         }
 
@@ -273,7 +272,7 @@
             metricColector.addStartupTimeMetricDuringVmRun(duringVmStartApp);
         }
 
-        shutdownMicrodroid(getDevice(), cid);
+        device.shutdownMicrodroid(microdroidDevice);
 
         // Run the app after the VM run and collect cold startup time.
         for (int i = 0; i < ROUND_COUNT; i++) {
diff --git a/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java b/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java
index 66cd211..ede838b 100644
--- a/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java
+++ b/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java
@@ -106,7 +106,7 @@
                 throws VirtualMachineException {
             VirtualMachine existingVm = mVmm.get(name);
             if (existingVm != null) {
-                existingVm.delete();
+                mVmm.delete(name);
             }
             return mVmm.create(name, config);
         }
diff --git a/tests/hostside/helper/java/com/android/microdroid/test/host/MicrodroidHostTestCaseBase.java b/tests/hostside/helper/java/com/android/microdroid/test/host/MicrodroidHostTestCaseBase.java
index 1a1dbeb..43fe615 100644
--- a/tests/hostside/helper/java/com/android/microdroid/test/host/MicrodroidHostTestCaseBase.java
+++ b/tests/hostside/helper/java/com/android/microdroid/test/host/MicrodroidHostTestCaseBase.java
@@ -38,13 +38,7 @@
 
 import java.io.File;
 import java.io.FileNotFoundException;
-import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Optional;
-import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Executors;
-import java.util.regex.Matcher;
-import java.util.regex.Pattern;
 
 public abstract class MicrodroidHostTestCaseBase extends BaseHostJUnit4Test {
     protected static final String TEST_ROOT = "/data/local/tmp/virt/";
@@ -160,53 +154,12 @@
         return result.getStdout().trim();
     }
 
-    // Same as runOnMicrodroid, but keeps retrying on error for maximum attempts times
-    // Each attempt with timeoutMs
-    public static String runOnMicrodroidRetryingOnFailure(
-            long timeoutMs, int attempts, String... cmd) {
-        CommandResult result = RunUtil.getDefault()
-                .runTimedCmdRetry(timeoutMs, MICRODROID_COMMAND_RETRY_INTERVAL_MILLIS, attempts,
-                        "adb", "-s", MICRODROID_SERIAL, "shell", join(cmd));
-        assertWithMessage("Command `" + Arrays.toString(cmd) + "` has failed")
-                .about(command_results())
-                .that(result)
-                .isSuccess();
-        return result.getStdout().trim();
-    }
-
     public static CommandResult runOnMicrodroidForResult(String... cmd) {
         final long timeoutMs = 30000; // 30 sec. Microdroid is extremely slow on GCE-on-CF.
         return RunUtil.getDefault()
                 .runTimedCmd(timeoutMs, "adb", "-s", MICRODROID_SERIAL, "shell", join(cmd));
     }
 
-    public static void pullMicrodroidFile(String path, File target) {
-        final long timeoutMs = 30000; // 30 sec. Microdroid is extremely slow on GCE-on-CF.
-        CommandResult result =
-                RunUtil.getDefault()
-                        .runTimedCmd(
-                                timeoutMs,
-                                "adb",
-                                "-s",
-                                MICRODROID_SERIAL,
-                                "pull",
-                                path,
-                                target.getPath());
-        assertWithMessage("pulling " + path + " from microdroid")
-                .about(command_results())
-                .that(result)
-                .isSuccess();
-    }
-
-    // Asserts the command will fail on Microdroid.
-    public static void assertFailedOnMicrodroid(String... cmd) {
-        CommandResult result = runOnMicrodroidForResult(cmd);
-        assertWithMessage("Microdroid command `" + join(cmd) + "` did not fail expectedly")
-                .about(command_results())
-                .that(result)
-                .isFailed();
-    }
-
     private static String join(String... strs) {
         return String.join(" ", Arrays.asList(strs));
     }
@@ -240,36 +193,6 @@
         return pathLine.substring("package:".length());
     }
 
-    public static String startMicrodroid(
-            ITestDevice androidDevice,
-            IBuildInfo buildInfo,
-            String apkName,
-            String packageName,
-            String configPath,
-            boolean debug,
-            int memoryMib,
-            Optional<Integer> numCpus)
-            throws DeviceNotAvailableException {
-        return startMicrodroid(androidDevice, buildInfo, apkName, packageName, null, configPath,
-                debug, memoryMib, numCpus);
-    }
-
-    public static String startMicrodroid(
-            ITestDevice androidDevice,
-            IBuildInfo buildInfo,
-            String apkName,
-            String packageName,
-            String[] extraIdsigPaths,
-            String configPath,
-            boolean debug,
-            int memoryMib,
-            Optional<Integer> numCpus)
-            throws DeviceNotAvailableException {
-        return startMicrodroid(androidDevice, buildInfo, apkName, null, packageName,
-                extraIdsigPaths, configPath, debug,
-                memoryMib, numCpus);
-    }
-
     private static void forwardFileToLog(CommandRunner android, String path, String tag)
             throws DeviceNotAvailableException {
         android.runWithTimeout(
@@ -281,93 +204,6 @@
                         + " | sed \\'s/^/" + tag + ": /g\\''\""); // add tags in front of lines
     }
 
-    public static String startMicrodroid(
-            ITestDevice androidDevice,
-            IBuildInfo buildInfo,
-            String apkName,
-            String apkPath,
-            String packageName,
-            String[] extraIdsigPaths,
-            String configPath,
-            boolean debug,
-            int memoryMib,
-            Optional<Integer> numCpus)
-            throws DeviceNotAvailableException {
-        CommandRunner android = new CommandRunner(androidDevice);
-
-        // Install APK if necessary
-        if (apkName != null) {
-            File apkFile = findTestFile(buildInfo, apkName);
-            androidDevice.installPackage(apkFile, /* reinstall */ true);
-        }
-
-        if (apkPath == null) {
-            apkPath = getPathForPackage(androidDevice, packageName);
-        }
-
-        android.run("mkdir", "-p", TEST_ROOT);
-
-        // This file is not what we provide. It will be created by the vm tool.
-        final String outApkIdsigPath = TEST_ROOT + apkName + ".idsig";
-
-        final String instanceImg = TEST_ROOT + INSTANCE_IMG;
-        final String logPath = LOG_PATH;
-        final String consolePath = CONSOLE_PATH;
-        final String debugFlag = debug ? "--debug full" : "";
-
-        // Run the VM
-        ArrayList<String> args = new ArrayList<>(Arrays.asList(
-                VIRT_APEX + "bin/vm",
-                "run-app",
-                "--daemonize",
-                "--log " + logPath,
-                "--console " + consolePath,
-                "--mem " + memoryMib,
-                numCpus.isPresent() ? "--cpus " + numCpus.get() : "",
-                debugFlag,
-                apkPath,
-                outApkIdsigPath,
-                instanceImg,
-                configPath));
-        if (extraIdsigPaths != null) {
-            for (String path : extraIdsigPaths) {
-                args.add("--extra-idsig");
-                args.add(path);
-            }
-        }
-        String ret = android.run(args.toArray(new String[0]));
-
-        // Redirect log.txt and console.txt to logd using logwrapper
-        // Keep redirecting as long as the expecting maximum test time. When an adb
-        // command times out, it may trigger the device recovery process, which
-        // disconnect adb, which terminates any live adb commands. See an example at
-        // b/194974010#comment25.
-        ExecutorService executor = Executors.newFixedThreadPool(2);
-        executor.execute(
-                () -> {
-                    try {
-                        forwardFileToLog(android, logPath, "MicrodroidLog");
-                    } catch (Exception e) {
-                        // Consume
-                    }
-                });
-
-        executor.execute(
-                () -> {
-                    try {
-                        forwardFileToLog(android, consolePath, "MicrodroidConsole");
-                    } catch (Exception e) {
-                        // Consume
-                    }
-                });
-
-        // Retrieve the CID from the vm tool output
-        Pattern pattern = Pattern.compile("with CID (\\d+)");
-        Matcher matcher = pattern.matcher(ret);
-        assertWithMessage("Failed to find CID").that(matcher.find()).isTrue();
-        return matcher.group(1);
-    }
-
     public static void shutdownMicrodroid(ITestDevice androidDevice, String cid)
             throws DeviceNotAvailableException {
         CommandRunner android = new CommandRunner(androidDevice);
@@ -376,17 +212,6 @@
         android.run(VIRT_APEX + "bin/vm", "stop", cid);
     }
 
-    public static void rootMicrodroid() throws InterruptedException {
-        runOnHostRetryingOnFailure(MICRODROID_COMMAND_TIMEOUT_MILLIS,
-                MICRODROID_ADB_CONNECT_MAX_ATTEMPTS, "adb", "-s", MICRODROID_SERIAL, "root");
-        // adbd reboots after root. Some commands (including wait-for-device) following this fails
-        // with error: closed. Hence, we disconnect and re-connect to the device before returning.
-        runOnHostRetryingOnFailure(MICRODROID_COMMAND_TIMEOUT_MILLIS,
-                MICRODROID_ADB_CONNECT_MAX_ATTEMPTS, "adb", "disconnect", MICRODROID_SERIAL);
-        runOnHostRetryingOnFailure(MICRODROID_COMMAND_TIMEOUT_MILLIS,
-                MICRODROID_ADB_CONNECT_MAX_ATTEMPTS, "adb", "connect", MICRODROID_SERIAL);
-    }
-
     // Establish an adb connection to microdroid by letting Android forward the connection to
     // microdroid. Wait until the connection is established and microdroid is booted.
     public static void adbConnectToMicrodroid(ITestDevice androidDevice, String cid) {
diff --git a/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java b/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
index 33788ed..f9de77e 100644
--- a/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
+++ b/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
@@ -17,12 +17,14 @@
 package com.android.microdroid.test;
 
 import static com.android.microdroid.test.host.CommandResultSubject.command_results;
+import static com.android.tradefed.device.TestDevice.MicrodroidBuilder;
 import static com.android.tradefed.testtype.DeviceJUnit4ClassRunner.TestLogData;
 
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth.assertWithMessage;
 
 import static org.hamcrest.CoreMatchers.containsString;
+import static org.hamcrest.CoreMatchers.is;
 import static org.junit.Assert.assertThat;
 import static org.junit.Assume.assumeTrue;
 
@@ -38,6 +40,8 @@
 import com.android.os.AtomsProto;
 import com.android.os.StatsLog;
 import com.android.tradefed.device.DeviceNotAvailableException;
+import com.android.tradefed.device.ITestDevice;
+import com.android.tradefed.device.TestDevice;
 import com.android.tradefed.testtype.DeviceJUnit4ClassRunner;
 import com.android.tradefed.testtype.DeviceJUnit4ClassRunner.TestMetrics;
 import com.android.tradefed.util.CommandResult;
@@ -64,8 +68,8 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
-import java.util.Optional;
 import java.util.concurrent.Callable;
+import java.util.function.Function;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
@@ -81,12 +85,16 @@
     // Number of vCPUs for testing purpose
     private static final int NUM_VCPUS = 3;
 
+    private static final int BOOT_COMPLETE_TIMEOUT = 30000; // 30 seconds
+
     @Rule public TestLogData mTestLogs = new TestLogData();
     @Rule public TestName mTestName = new TestName();
     @Rule public TestMetrics mMetrics = new TestMetrics();
 
     private String mMetricPrefix;
 
+    private ITestDevice mMicrodroidDevice;
+
     private int minMemorySize() throws DeviceNotAvailableException {
         CommandRunner android = new CommandRunner(getDevice());
         String abi = android.run("getprop", "ro.product.cpu.abi");
@@ -99,10 +107,6 @@
         throw new AssertionError("Unsupported ABI: " + abi);
     }
 
-    private void waitForBootComplete() {
-        runOnMicrodroidForResult("watch -e \"getprop dev.bootcomplete | grep '^0$'\"");
-    }
-
     private static JSONObject newPartition(String label, String path) {
         return new JSONObject(Map.of("label", label, "path", path));
     }
@@ -447,16 +451,18 @@
     private boolean isTombstoneGeneratedWithConfig(String configPath) throws Exception {
         // Note this test relies on logcat values being printed by tombstone_transmit on
         // and the reeceiver on host (virtualization_service)
-        final String cid =
-                startMicrodroid(
-                        getDevice(),
-                        getBuild(),
-                        APK_NAME,
-                        PACKAGE_NAME,
-                        configPath,
-                        /* debug */ true,
-                        minMemorySize(),
-                        Optional.of(NUM_VCPUS));
+        mMicrodroidDevice = MicrodroidBuilder
+                .fromDevicePath(getPathForPackage(PACKAGE_NAME), configPath)
+                .debugLevel("full")
+                .memoryMib(minMemorySize())
+                .numCpus(NUM_VCPUS)
+                .build((TestDevice) getDevice());
+        mMicrodroidDevice.waitForBootComplete(BOOT_COMPLETE_TIMEOUT);
+        mMicrodroidDevice.enableAdbRoot();
+
+        CommandRunner microdroid = new CommandRunner(mMicrodroidDevice);
+        microdroid.run("kill", "-SIGSEGV", "$(pidof microdroid_launcher)");
+
         // check until microdroid is shut down
         CommandRunner android = new CommandRunner(getDevice());
         android.runWithTimeout(15000, "logcat", "-m", "1", "-e", "'crosvm has exited normally'");
@@ -485,7 +491,7 @@
     }
 
     @Test
-    public void testTelemetryPushedAtomsOfEventMetrics() throws Exception {
+    public void testTelemetryPushedAtoms() throws Exception {
         // Reset statsd config and report before the test
         ConfigUtils.removeConfig(getDevice());
         ReportUtils.clearReports(getDevice());
@@ -499,22 +505,28 @@
         ConfigUtils.uploadConfigForPushedAtoms(getDevice(), PACKAGE_NAME, atomIds);
 
         // Create VM with microdroid
+        TestDevice device = (TestDevice) getDevice();
         final String configPath = "assets/vm_config_apex.json"; // path inside the APK
-        final String cid =
-                startMicrodroid(
-                        getDevice(),
-                        getBuild(),
-                        APK_NAME,
-                        PACKAGE_NAME,
-                        configPath,
-                        /* debug */ true,
-                        minMemorySize(),
-                        Optional.of(NUM_VCPUS));
+        ITestDevice microdroid =
+                MicrodroidBuilder.fromDevicePath(getPathForPackage(PACKAGE_NAME), configPath)
+                        .debugLevel("full")
+                        .memoryMib(minMemorySize())
+                        .numCpus(NUM_VCPUS)
+                        .build(device);
+        microdroid.waitForBootComplete(BOOT_COMPLETE_TIMEOUT);
+        device.shutdownMicrodroid(microdroid);
 
-        // Check VmCreationRequested atom and clear the statsd report
-        List<StatsLog.EventMetricData> data;
-        data = ReportUtils.getEventMetricDataList(getDevice());
-        assertThat(data).hasSize(1);
+        List<StatsLog.EventMetricData> data = new ArrayList<>();
+        assertThatEventually(
+                10000,
+                () -> {
+                    data.addAll(ReportUtils.getEventMetricDataList(getDevice()));
+                    return data.size();
+                },
+                is(3)
+        );
+
+        // Check VmCreationRequested atom
         assertThat(data.get(0).getAtom().getPushedCase().getNumber()).isEqualTo(
                 AtomsProto.Atom.VM_CREATION_REQUESTED_FIELD_NUMBER);
         AtomsProto.VmCreationRequested atomVmCreationRequested =
@@ -532,29 +544,16 @@
         assertThat(atomVmCreationRequested.getApexes())
                 .isEqualTo("com.android.art:com.android.compos:com.android.sdkext");
 
-        // Boot VM with microdroid
-        adbConnectToMicrodroid(getDevice(), cid);
-        waitForBootComplete();
-
-        // Check VmBooted atom and clear the statsd report
-        data = ReportUtils.getEventMetricDataList(getDevice());
-        assertThat(data).hasSize(1);
-        assertThat(data.get(0).getAtom().getPushedCase().getNumber())
+        // Check VmBooted atom
+        assertThat(data.get(1).getAtom().getPushedCase().getNumber())
                 .isEqualTo(AtomsProto.Atom.VM_BOOTED_FIELD_NUMBER);
-        AtomsProto.VmBooted atomVmBooted = data.get(0).getAtom().getVmBooted();
+        AtomsProto.VmBooted atomVmBooted = data.get(1).getAtom().getVmBooted();
         assertThat(atomVmBooted.getVmIdentifier()).isEqualTo("VmRunApp");
 
-        // Shutdown VM with microdroid
-        shutdownMicrodroid(getDevice(), cid);
-        // TODO: make sure the VM is completely shut down while 'vm stop' command running.
-        Thread.sleep(1000);
-
-        // Check VmExited atom and clear the statsd report
-        data = ReportUtils.getEventMetricDataList(getDevice());
-        assertThat(data).hasSize(1);
-        assertThat(data.get(0).getAtom().getPushedCase().getNumber())
+        // Check VmExited atom
+        assertThat(data.get(2).getAtom().getPushedCase().getNumber())
                 .isEqualTo(AtomsProto.Atom.VM_EXITED_FIELD_NUMBER);
-        AtomsProto.VmExited atomVmExited = data.get(0).getAtom().getVmExited();
+        AtomsProto.VmExited atomVmExited = data.get(2).getAtom().getVmExited();
         assertThat(atomVmExited.getVmIdentifier()).isEqualTo("VmRunApp");
         assertThat(atomVmExited.getDeathReason()).isEqualTo(AtomsProto.VmExited.DeathReason.KILLED);
 
@@ -566,100 +565,52 @@
     }
 
     @Test
-    public void testTelemetryPushedAtomsOfValueMetrics() throws Exception {
-        // Reset statsd config and report before the test
-        ConfigUtils.removeConfig(getDevice());
-        ReportUtils.clearReports(getDevice());
-
-        // Setup statsd config
-        int[] atomIds = {
-            AtomsProto.Atom.VM_CPU_STATUS_REPORTED_FIELD_NUMBER,
-            AtomsProto.Atom.VM_MEM_STATUS_REPORTED_FIELD_NUMBER,
-        };
-        ConfigUtils.uploadConfigForPushedAtoms(getDevice(), PACKAGE_NAME, atomIds);
-
-        // Create VM with microdroid
-        final String configPath = "assets/vm_config_apex.json"; // path inside the APK
-        final String cid =
-                startMicrodroid(
-                        getDevice(),
-                        getBuild(),
-                        APK_NAME,
-                        PACKAGE_NAME,
-                        configPath,
-                        /* debug */ true,
-                        minMemorySize(),
-                        Optional.of(NUM_VCPUS));
-
-        // Boot VM with microdroid
-        adbConnectToMicrodroid(getDevice(), cid);
-        waitForBootComplete();
-
-        // Check VmCpuStatusReported and VmMemStatusReported atoms and clear the statsd report
-        List<StatsLog.EventMetricData> data;
-        data = ReportUtils.getEventMetricDataList(getDevice());
-        assertThat(data.size() >= 2).isTrue();
-        assertThat(data.get(0).getAtom().getPushedCase().getNumber())
-                .isEqualTo(AtomsProto.Atom.VM_CPU_STATUS_REPORTED_FIELD_NUMBER);
-        assertThat(data.get(1).getAtom().getPushedCase().getNumber())
-                .isEqualTo(AtomsProto.Atom.VM_MEM_STATUS_REPORTED_FIELD_NUMBER);
-
-        // Shutdown VM with microdroid
-        shutdownMicrodroid(getDevice(), cid);
-    }
-
-    @Test
     @CddTest(requirements = {"9.17/C-1-1", "9.17/C-1-2", "9.17/C/1-3"})
     public void testMicrodroidBoots() throws Exception {
         final String configPath = "assets/vm_config.json"; // path inside the APK
-        final String cid =
-                startMicrodroid(
-                        getDevice(),
-                        getBuild(),
-                        APK_NAME,
-                        PACKAGE_NAME,
-                        configPath,
-                        /* debug */ true,
-                        minMemorySize(),
-                        Optional.of(NUM_VCPUS));
-        adbConnectToMicrodroid(getDevice(), cid);
-        waitForBootComplete();
+        mMicrodroidDevice =
+                MicrodroidBuilder.fromDevicePath(getPathForPackage(PACKAGE_NAME), configPath)
+                        .debugLevel("full")
+                        .memoryMib(minMemorySize())
+                        .numCpus(NUM_VCPUS)
+                        .build((TestDevice) getDevice());
+        mMicrodroidDevice.waitForBootComplete(BOOT_COMPLETE_TIMEOUT);
+        CommandRunner microdroid = new CommandRunner(mMicrodroidDevice);
+
         // Test writing to /data partition
-        runOnMicrodroid("echo MicrodroidTest > /data/local/tmp/test.txt");
-        assertThat(runOnMicrodroid("cat /data/local/tmp/test.txt")).isEqualTo("MicrodroidTest");
+        microdroid.run("echo MicrodroidTest > /data/local/tmp/test.txt");
+        assertThat(microdroid.run("cat /data/local/tmp/test.txt")).isEqualTo("MicrodroidTest");
 
         // Check if the APK & its idsig partitions exist
         final String apkPartition = "/dev/block/by-name/microdroid-apk";
-        assertThat(runOnMicrodroid("ls", apkPartition)).isEqualTo(apkPartition);
+        assertThat(microdroid.run("ls", apkPartition)).isEqualTo(apkPartition);
         final String apkIdsigPartition = "/dev/block/by-name/microdroid-apk-idsig";
-        assertThat(runOnMicrodroid("ls", apkIdsigPartition)).isEqualTo(apkIdsigPartition);
+        assertThat(microdroid.run("ls", apkIdsigPartition)).isEqualTo(apkIdsigPartition);
         // Check the vm-instance partition as well
         final String vmInstancePartition = "/dev/block/by-name/vm-instance";
-        assertThat(runOnMicrodroid("ls", vmInstancePartition)).isEqualTo(vmInstancePartition);
+        assertThat(microdroid.run("ls", vmInstancePartition)).isEqualTo(vmInstancePartition);
 
         // Check if the native library in the APK is has correct filesystem info
-        final String[] abis = runOnMicrodroid("getprop", "ro.product.cpu.abilist").split(",");
+        final String[] abis = microdroid.run("getprop", "ro.product.cpu.abilist").split(",");
         assertThat(abis).hasLength(1);
         final String testLib = "/mnt/apk/lib/" + abis[0] + "/MicrodroidTestNativeLib.so";
         final String label = "u:object_r:system_file:s0";
-        assertThat(runOnMicrodroid("ls", "-Z", testLib)).isEqualTo(label + " " + testLib);
+        assertThat(microdroid.run("ls", "-Z", testLib)).isEqualTo(label + " " + testLib);
 
         // Check that no denials have happened so far
         CommandRunner android = new CommandRunner(getDevice());
         assertThat(android.tryRun("egrep", "'avc:[[:space:]]{1,2}denied'", LOG_PATH)).isNull();
 
-        assertThat(runOnMicrodroid("cat /proc/cpuinfo | grep processor | wc -l"))
+        assertThat(microdroid.run("cat /proc/cpuinfo | grep processor | wc -l"))
                 .isEqualTo(Integer.toString(NUM_VCPUS));
 
         // Check that selinux is enabled
-        assertThat(runOnMicrodroid("getenforce")).isEqualTo("Enforcing");
+        assertThat(microdroid.run("getenforce")).isEqualTo("Enforcing");
 
         // TODO(b/176805428): adb is broken for nested VM
         if (!isCuttlefish()) {
             // Check neverallow rules on microdroid
-            File policyFile = FileUtil.createTempFile("microdroid_sepolicy", "");
-            pullMicrodroidFile("/sys/fs/selinux/policy", policyFile);
-
+            File policyFile = mMicrodroidDevice.pullFile("/sys/fs/selinux/policy");
             File generalPolicyConfFile = findTestFile("microdroid_general_sepolicy.conf");
             File sepolicyAnalyzeBin = findTestFile("sepolicy-analyze");
 
@@ -678,38 +629,41 @@
                     .that(result)
                     .isSuccess();
         }
-
-        shutdownMicrodroid(getDevice(), cid);
     }
 
     @Test
     public void testMicrodroidRamUsage() throws Exception {
         final String configPath = "assets/vm_config.json";
-        final String cid =
-                startMicrodroid(
-                        getDevice(),
-                        getBuild(),
-                        APK_NAME,
-                        PACKAGE_NAME,
-                        configPath,
-                        /* debug */ true,
-                        minMemorySize(),
-                        Optional.of(NUM_VCPUS));
-        adbConnectToMicrodroid(getDevice(), cid);
-        waitForBootComplete();
-        rootMicrodroid();
+        mMicrodroidDevice = MicrodroidBuilder
+                .fromDevicePath(getPathForPackage(PACKAGE_NAME), configPath)
+                .debugLevel("full")
+                .memoryMib(minMemorySize())
+                .numCpus(NUM_VCPUS)
+                .build((TestDevice) getDevice());
+        mMicrodroidDevice.waitForBootComplete(BOOT_COMPLETE_TIMEOUT);
+        mMicrodroidDevice.enableAdbRoot();
+
+        CommandRunner microdroid = new CommandRunner(mMicrodroidDevice);
+        Function<String, String> microdroidExec =
+                (cmd) -> {
+                    try {
+                        return microdroid.run(cmd);
+                    } catch (Exception ex) {
+                        throw new IllegalStateException(ex);
+                    }
+                };
 
         for (Map.Entry<String, Long> stat :
-                ProcessUtil.getProcessMemoryMap(cmd -> runOnMicrodroid(cmd)).entrySet()) {
+                ProcessUtil.getProcessMemoryMap(microdroidExec).entrySet()) {
             mMetrics.addTestMetric(
                     mMetricPrefix + "meminfo/" + stat.getKey().toLowerCase(),
                     stat.getValue().toString());
         }
 
         for (Map.Entry<Integer, String> proc :
-                ProcessUtil.getProcessMap(cmd -> runOnMicrodroid(cmd)).entrySet()) {
+                ProcessUtil.getProcessMap(microdroidExec).entrySet()) {
             for (Map.Entry<String, Long> stat :
-                    ProcessUtil.getProcessSmapsRollup(proc.getKey(), cmd -> runOnMicrodroid(cmd))
+                    ProcessUtil.getProcessSmapsRollup(proc.getKey(), microdroidExec)
                             .entrySet()) {
                 String name = stat.getKey().toLowerCase();
                 mMetrics.addTestMetric(
@@ -717,8 +671,6 @@
                         stat.getValue().toString());
             }
         }
-
-        shutdownMicrodroid(getDevice(), cid);
     }
 
     @Test
@@ -760,6 +712,7 @@
     public void setUp() throws Exception {
         testIfDeviceIsCapable(getDevice());
         mMetricPrefix = getMetricPrefix() + "microdroid/";
+        mMicrodroidDevice = null;
 
         prepareVirtualizationTestSetup(getDevice());
 
@@ -771,6 +724,10 @@
 
     @After
     public void shutdown() throws Exception {
+        if (mMicrodroidDevice != null) {
+            ((TestDevice) getDevice()).shutdownMicrodroid(mMicrodroidDevice);
+        }
+
         cleanUpVirtualizationTestSetup(getDevice());
 
         archiveLogThenDelete(
diff --git a/tests/testapk/Android.bp b/tests/testapk/Android.bp
index 42abbbf..8972046 100644
--- a/tests/testapk/Android.bp
+++ b/tests/testapk/Android.bp
@@ -22,7 +22,7 @@
     libs: ["android.system.virtualmachine"],
     jni_libs: [
         "MicrodroidTestNativeLib",
-        "MicrodroidTestNativeCrashLib",
+        "MicrodroidIdleNativeLib",
     ],
     platform_apis: true,
     use_embedded_native_libs: true,
@@ -35,6 +35,7 @@
     name: "MicrodroidTestNativeLib",
     srcs: ["src/native/testbinary.cpp"],
     stl: "libc++_static",
+    header_libs: ["vm_payload_restricted_headers"],
     shared_libs: [
         "libbinder_ndk",
         "MicrodroidTestNativeLibSub",
@@ -50,13 +51,6 @@
 }
 
 cc_library_shared {
-    name: "MicrodroidTestNativeCrashLib",
-    header_libs: ["vm_payload_headers"],
-    srcs: ["src/native/crashbinary.cpp"],
-    stl: "libc++_static",
-}
-
-cc_library_shared {
     name: "MicrodroidTestNativeLibSub",
     srcs: ["src/native/testlib.cpp"],
     stl: "libc++_static",
diff --git a/tests/testapk/assets/vm_config_crash.json b/tests/testapk/assets/vm_config_crash.json
index 2951fdf..3ec34a3 100644
--- a/tests/testapk/assets/vm_config_crash.json
+++ b/tests/testapk/assets/vm_config_crash.json
@@ -4,7 +4,7 @@
     },
     "task": {
       "type": "microdroid_launcher",
-      "command": "MicrodroidTestNativeCrashLib.so"
+      "command": "MicrodroidIdleNativeLib.so"
     },
     "export_tombstones": true
   }
diff --git a/tests/testapk/assets/vm_config_crash_no_tombstone.json b/tests/testapk/assets/vm_config_crash_no_tombstone.json
index 583f87b..9678e38 100644
--- a/tests/testapk/assets/vm_config_crash_no_tombstone.json
+++ b/tests/testapk/assets/vm_config_crash_no_tombstone.json
@@ -4,7 +4,7 @@
     },
     "task": {
       "type": "microdroid_launcher",
-      "command": "MicrodroidTestNativeCrashLib.so"
+      "command": "MicrodroidIdleNativeLib.so"
     },
     "export_tombstones": false
   }
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 e5052bf..cc623a8 100644
--- a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
+++ b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
@@ -25,14 +25,21 @@
 
 import static java.nio.file.StandardCopyOption.REPLACE_EXISTING;
 
+import android.content.Context;
 import android.os.Build;
 import android.os.ParcelFileDescriptor;
+import android.os.ServiceSpecificException;
 import android.os.SystemProperties;
+import android.system.virtualmachine.ParcelVirtualMachine;
 import android.system.virtualmachine.VirtualMachine;
 import android.system.virtualmachine.VirtualMachineCallback;
 import android.system.virtualmachine.VirtualMachineConfig;
+import android.system.virtualmachine.VirtualMachineException;
+import android.system.virtualmachine.VirtualMachineManager;
 import android.util.Log;
 
+import androidx.test.core.app.ApplicationProvider;
+
 import com.android.compatibility.common.util.CddTest;
 import com.android.microdroid.test.device.MicrodroidDeviceTestBase;
 import com.android.microdroid.testservice.ITestService;
@@ -52,6 +59,8 @@
 import java.io.IOException;
 import java.io.RandomAccessFile;
 import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
 import java.util.List;
 import java.util.OptionalLong;
 import java.util.UUID;
@@ -104,7 +113,7 @@
                 .setPayloadBinaryPath("MicrodroidTestNativeLib.so")
                 .setMemoryMib(minMemoryRequired())
                 .build();
-        VirtualMachine vm = mInner.forceCreateNewVirtualMachine("test_vm_extra_apk", config);
+        VirtualMachine vm = mInner.forceCreateNewVirtualMachine("test_vm", config);
 
         TestResults testResults = runVmTestService(vm);
         assertThat(testResults.mException).isNull();
@@ -162,6 +171,33 @@
     @CddTest(requirements = {
             "9.17/C-1-1",
     })
+    public void deleteVm() throws Exception {
+        assumeSupportedKernel();
+
+        VirtualMachineConfig config = mInner.newVmConfigBuilder()
+                .setPayloadBinaryPath("MicrodroidTestNativeLib.so")
+                .setMemoryMib(minMemoryRequired())
+                .build();
+
+        VirtualMachine vm = mInner.forceCreateNewVirtualMachine("test_vm_delete",
+                config);
+        VirtualMachineManager vmm = mInner.getVirtualMachineManager();
+        vmm.delete("test_vm_delete");
+
+        // VM should no longer exist
+        assertThat(vmm.get("test_vm_delete")).isNull();
+
+        // Can't start the VM even with an existing reference
+        assertThrows(VirtualMachineException.class, vm::run);
+
+        // Can't delete the VM since it no longer exists
+        assertThrows(VirtualMachineException.class, () -> vmm.delete("test_vm_delete"));
+    }
+
+    @Test
+    @CddTest(requirements = {
+            "9.17/C-1-1",
+    })
     public void validApkPathIsAccepted() throws Exception {
         assumeSupportedKernel();
 
@@ -260,8 +296,7 @@
         // Try to run the VM again with the previous instance.img
         // We need to make sure that no changes on config don't invalidate the identity, to compare
         // the result with the below "different debug level" test.
-        File vmRoot = new File(getContext().getFilesDir(), "vm");
-        File vmInstance = new File(new File(vmRoot, "test_vm"), "instance.img");
+        File vmInstance = getVmFile("test_vm", "instance.img");
         File vmInstanceBackup = File.createTempFile("instance", ".img");
         Files.copy(vmInstance.toPath(), vmInstanceBackup.toPath(), REPLACE_EXISTING);
         mInner.forceCreateNewVirtualMachine("test_vm", normalConfig);
@@ -303,7 +338,10 @@
                     }
                 };
         listener.runToFinish(TAG, vm);
-        assertThat(exception.getNow(null)).isNull();
+        Exception e = exception.getNow(null);
+        if (e != null) {
+            throw e;
+        }
         return vmCdis;
     }
 
@@ -404,6 +442,24 @@
         }
     }
 
+    @Test
+    @CddTest(requirements = {
+            "9.17/C-1-1",
+            "9.17/C-1-2"
+    })
+    public void accessToCdisIsRestricted() throws Exception {
+        assumeSupportedKernel();
+
+        VirtualMachineConfig config = mInner.newVmConfigBuilder()
+                .setPayloadBinaryPath("MicrodroidTestNativeLib.so")
+                .setDebugLevel(DEBUG_LEVEL_FULL)
+                .build();
+        mInner.forceCreateNewVirtualMachine("test_vm", config);
+
+        assertThrows(ServiceSpecificException.class, () -> launchVmAndGetCdis("test_vm"));
+    }
+
+
     private static final UUID MICRODROID_PARTITION_UUID =
             UUID.fromString("cf9afe9a-0662-11ec-a329-c32663a09d75");
     private static final UUID U_BOOT_AVB_PARTITION_UUID =
@@ -447,10 +503,7 @@
 
         mInner.forceCreateNewVirtualMachine(vmName, config);
         assertThat(tryBootVm(TAG, vmName).payloadStarted).isTrue();
-
-        File vmRoot = new File(getContext().getFilesDir(), "vm");
-        File vmDir = new File(vmRoot, vmName);
-        File instanceImgPath = new File(vmDir, "instance.img");
+        File instanceImgPath = getVmFile(vmName, "instance.img");
         return new RandomAccessFile(instanceImgPath, "rw");
     }
 
@@ -546,6 +599,41 @@
         assertThat(vm).isNotEqualTo(newVm);
     }
 
+    @Test
+    public void vmConvertsToValidParcelVm() throws Exception {
+        // Arrange
+        VirtualMachineConfig config =
+                mInner.newVmConfigBuilder()
+                        .setPayloadBinaryPath("MicrodroidTestNativeLib.so")
+                        .setDebugLevel(DEBUG_LEVEL_NONE)
+                        .build();
+        String vmName = "test_vm";
+        VirtualMachine vm = mInner.forceCreateNewVirtualMachine(vmName, config);
+
+        // Action
+        ParcelVirtualMachine parcelVm = vm.toParcelVirtualMachine();
+
+        // Asserts
+        assertFileContentsAreEqual(parcelVm.getConfigFd(), vmName, "config.xml");
+        assertFileContentsAreEqual(parcelVm.getInstanceImgFd(), vmName, "instance.img");
+    }
+
+    private void assertFileContentsAreEqual(
+            ParcelFileDescriptor parcelFd, String vmName, String fileName) throws IOException {
+        File file = getVmFile(vmName, fileName);
+        // Use try-with-resources to close the files automatically after assert.
+        try (FileInputStream input1 = new FileInputStream(parcelFd.getFileDescriptor());
+                FileInputStream input2 = new FileInputStream(file)) {
+            assertThat(input1.readAllBytes()).isEqualTo(input2.readAllBytes());
+        }
+    }
+
+    private File getVmFile(String vmName, String fileName) {
+        Context context = ApplicationProvider.getApplicationContext();
+        Path filePath = Paths.get(context.getDataDir().getPath(), "vm", vmName, fileName);
+        return filePath.toFile();
+    }
+
     private int minMemoryRequired() {
         if (Build.SUPPORTED_ABIS.length > 0) {
             String primaryAbi = Build.SUPPORTED_ABIS[0];
diff --git a/tests/testapk/src/native/crashbinary.cpp b/tests/testapk/src/native/crashbinary.cpp
deleted file mode 100644
index a0edc40..0000000
--- a/tests/testapk/src/native/crashbinary.cpp
+++ /dev/null
@@ -1,26 +0,0 @@
-/*
- * Copyright 2022 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.
- */
-
-#include <stdio.h>
-#include <stdlib.h>
-
-#include "vm_main.h"
-
-// A VM payload that crashes as soon as it starts, to allow us to exercise that error path.
-extern "C" int AVmPayload_main() {
-    printf("test crash!!!!\n");
-    abort();
-}
diff --git a/tests/testapk/src/native/testbinary.cpp b/tests/testapk/src/native/testbinary.cpp
index 1a3e940..48942dc 100644
--- a/tests/testapk/src/native/testbinary.cpp
+++ b/tests/testapk/src/native/testbinary.cpp
@@ -28,7 +28,7 @@
 #include <sys/system_properties.h>
 #include <unistd.h>
 #include <vm_main.h>
-#include <vm_payload.h>
+#include <vm_payload_restricted.h>
 
 #include <string>
 
@@ -79,7 +79,7 @@
             if (!AVmPayload_getVmInstanceSecret(identifier, sizeof(identifier), out->data(),
                                                 out->size())) {
                 return ndk::ScopedAStatus::
-                        fromServiceSpecificErrorWithMessage(0, "Failed to VM instance secret");
+                        fromServiceSpecificErrorWithMessage(0, "Failed to get VM instance secret");
             }
             return ndk::ScopedAStatus::ok();
         }
diff --git a/virtualizationservice/Android.bp b/virtualizationservice/Android.bp
index 26d41c9..d6f4607 100644
--- a/virtualizationservice/Android.bp
+++ b/virtualizationservice/Android.bp
@@ -71,5 +71,8 @@
 rust_test {
     name: "virtualizationservice_device_test",
     defaults: ["virtualizationservice_defaults"],
+    rustlibs: [
+        "libtempfile",
+    ],
     test_suites: ["general-tests"],
 }
diff --git a/virtualizationservice/aidl/android/system/virtualmachineservice/IVirtualMachineService.aidl b/virtualizationservice/aidl/android/system/virtualmachineservice/IVirtualMachineService.aidl
index 4fa5fa0..e8c1724 100644
--- a/virtualizationservice/aidl/android/system/virtualmachineservice/IVirtualMachineService.aidl
+++ b/virtualizationservice/aidl/android/system/virtualmachineservice/IVirtualMachineService.aidl
@@ -16,8 +16,6 @@
 package android.system.virtualmachineservice;
 
 import android.system.virtualizationcommon.ErrorCode;
-import android.system.virtualmachineservice.VirtualMachineCpuStatus;
-import android.system.virtualmachineservice.VirtualMachineMemStatus;
 
 /** {@hide} */
 interface IVirtualMachineService {
@@ -58,14 +56,4 @@
      * Notifies that an error has occurred inside the VM..
      */
     void notifyError(ErrorCode errorCode, in String message);
-
-    /**
-     * Notifies the current CPU status of the VM.
-     */
-    void notifyCpuStatus(in VirtualMachineCpuStatus cpuStatus);
-
-    /**
-     * Notifies the current memory status of the VM.
-     */
-    void notifyMemStatus(in VirtualMachineMemStatus memStatus);
 }
diff --git a/virtualizationservice/aidl/android/system/virtualmachineservice/VirtualMachineCpuStatus.aidl b/virtualizationservice/aidl/android/system/virtualmachineservice/VirtualMachineCpuStatus.aidl
deleted file mode 100644
index 307c3f9..0000000
--- a/virtualizationservice/aidl/android/system/virtualmachineservice/VirtualMachineCpuStatus.aidl
+++ /dev/null
@@ -1,23 +0,0 @@
-/*
- * Copyright 2022 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 android.system.virtualmachineservice;
-
-parcelable VirtualMachineCpuStatus {
-    long cpu_time_user;
-    long cpu_time_nice;
-    long cpu_time_sys;
-    long cpu_time_idle;
-}
diff --git a/virtualizationservice/aidl/android/system/virtualmachineservice/VirtualMachineMemStatus.aidl b/virtualizationservice/aidl/android/system/virtualmachineservice/VirtualMachineMemStatus.aidl
deleted file mode 100644
index 3de57c6..0000000
--- a/virtualizationservice/aidl/android/system/virtualmachineservice/VirtualMachineMemStatus.aidl
+++ /dev/null
@@ -1,24 +0,0 @@
-/*
- * Copyright 2022 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 android.system.virtualmachineservice;
-
-parcelable VirtualMachineMemStatus {
-    long mem_total;
-    long mem_free;
-    long mem_available;
-    long mem_buffer;
-    long mem_cached;
-}
diff --git a/virtualizationservice/src/aidl.rs b/virtualizationservice/src/aidl.rs
index b4ce9d2..bc697e3 100644
--- a/virtualizationservice/src/aidl.rs
+++ b/virtualizationservice/src/aidl.rs
@@ -14,10 +14,7 @@
 
 //! Implementation of the AIDL interface of the VirtualizationService.
 
-use crate::atom::{
-    write_vm_booted_stats, write_vm_cpu_status_stats, write_vm_creation_stats,
-    write_vm_mem_status_stats,
-};
+use crate::atom::{write_vm_booted_stats, write_vm_creation_stats};
 use crate::composite::make_composite_image;
 use crate::crosvm::{CrosvmConfig, DiskFile, PayloadState, VmInstance, VmState};
 use crate::payload::{add_microdroid_payload_images, add_microdroid_system_images};
@@ -39,13 +36,9 @@
     VirtualMachineRawConfig::VirtualMachineRawConfig,
     VirtualMachineState::VirtualMachineState,
 };
-use android_system_virtualmachineservice::aidl::android::system::virtualmachineservice::{
-    IVirtualMachineService::{
+use android_system_virtualmachineservice::aidl::android::system::virtualmachineservice::IVirtualMachineService::{
         BnVirtualMachineService, IVirtualMachineService, VM_BINDER_SERVICE_PORT,
         VM_STREAM_SERVICE_PORT, VM_TOMBSTONES_SERVICE_PORT,
-    },
-    VirtualMachineCpuStatus::VirtualMachineCpuStatus,
-    VirtualMachineMemStatus::VirtualMachineMemStatus,
 };
 use anyhow::{anyhow, bail, Context, Result};
 use apkverify::{HashAlgorithm, V4Signature};
@@ -1147,36 +1140,6 @@
             ))
         }
     }
-
-    fn notifyCpuStatus(&self, status: &VirtualMachineCpuStatus) -> binder::Result<()> {
-        let cid = self.cid;
-        if let Some(vm) = self.state.lock().unwrap().get_vm(cid) {
-            info!("VM with CID {} reported its CPU status", cid);
-            write_vm_cpu_status_stats(vm.requester_uid as i32, &vm.name, status);
-            Ok(())
-        } else {
-            error!("notifyCurrentStatus is called from an unknown CID {}", cid);
-            Err(Status::new_service_specific_error_str(
-                -1,
-                Some(format!("cannot find a VM with CID {}", cid)),
-            ))
-        }
-    }
-
-    fn notifyMemStatus(&self, status: &VirtualMachineMemStatus) -> binder::Result<()> {
-        let cid = self.cid;
-        if let Some(vm) = self.state.lock().unwrap().get_vm(cid) {
-            info!("VM with CID {} reported its memory status", cid);
-            write_vm_mem_status_stats(vm.requester_uid as i32, &vm.name, status);
-            Ok(())
-        } else {
-            error!("notifyCurrentStatus is called from an unknown CID {}", cid);
-            Err(Status::new_service_specific_error_str(
-                -1,
-                Some(format!("cannot find a VM with CID {}", cid)),
-            ))
-        }
-    }
 }
 
 impl VirtualMachineService {
diff --git a/virtualizationservice/src/atom.rs b/virtualizationservice/src/atom.rs
index 8c46ac5..20f88e7 100644
--- a/virtualizationservice/src/atom.rs
+++ b/virtualizationservice/src/atom.rs
@@ -22,17 +22,13 @@
     VirtualMachineConfig::VirtualMachineConfig,
 };
 use android_system_virtualizationservice::binder::{Status, Strong};
-use android_system_virtualmachineservice::aidl::android::system::virtualmachineservice::{
-    VirtualMachineCpuStatus::VirtualMachineCpuStatus,
-    VirtualMachineMemStatus::VirtualMachineMemStatus,
-};
 use anyhow::{anyhow, Result};
 use binder::{ParcelFileDescriptor, ThreadState};
 use log::{trace, warn};
 use microdroid_payload_config::VmPayloadConfig;
-use statslog_virtualization_rust::{
-    vm_booted, vm_cpu_status_reported, vm_creation_requested, vm_exited, vm_mem_status_reported,
-};
+use rustutils::system_properties;
+use statslog_virtualization_rust::{vm_booted, vm_creation_requested, vm_exited};
+use std::thread;
 use std::time::{Duration, SystemTime};
 use zip::ZipArchive;
 
@@ -92,17 +88,16 @@
             binder_exception_code = e.exception_code() as i32;
         }
     }
-
     let (vm_identifier, config_type, num_cpus, memory_mib, apexes) = match config {
         VirtualMachineConfig::AppConfig(config) => (
-            &config.name,
+            config.name.clone(),
             vm_creation_requested::ConfigType::VirtualMachineAppConfig,
             config.numCpus,
             config.memoryMib,
             get_apex_list(config),
         ),
         VirtualMachineConfig::RawConfig(config) => (
-            &config.name,
+            config.name.clone(),
             vm_creation_requested::ConfigType::VirtualMachineRawConfig,
             config.numCpus,
             config.memoryMib,
@@ -110,150 +105,139 @@
         ),
     };
 
-    let vm_creation_requested = vm_creation_requested::VmCreationRequested {
-        uid: ThreadState::get_calling_uid() as i32,
-        vm_identifier,
-        hypervisor: vm_creation_requested::Hypervisor::Pkvm,
-        is_protected,
-        creation_succeeded,
-        binder_exception_code,
-        config_type,
-        num_cpus,
-        cpu_affinity: "", // deprecated
-        memory_mib,
-        apexes: &apexes,
-        // TODO(seungjaeyoo) Fill information about task_profile
-        // TODO(seungjaeyoo) Fill information about disk_image for raw config
-    };
+    let uid = ThreadState::get_calling_uid() as i32;
+    thread::spawn(move || {
+        let vm_creation_requested = vm_creation_requested::VmCreationRequested {
+            uid,
+            vm_identifier: &vm_identifier,
+            hypervisor: vm_creation_requested::Hypervisor::Pkvm,
+            is_protected,
+            creation_succeeded,
+            binder_exception_code,
+            config_type,
+            num_cpus,
+            cpu_affinity: "", // deprecated
+            memory_mib,
+            apexes: &apexes,
+            // TODO(seungjaeyoo) Fill information about task_profile
+            // TODO(seungjaeyoo) Fill information about disk_image for raw config
+        };
 
-    match vm_creation_requested.stats_write() {
-        Err(e) => {
-            warn!("statslog_rust failed with error: {}", e);
+        wait_for_statsd().unwrap_or_else(|e| warn!("failed to wait for statsd with error: {}", e));
+        match vm_creation_requested.stats_write() {
+            Err(e) => {
+                warn!("statslog_rust failed with error: {}", e);
+            }
+            Ok(_) => trace!("statslog_rust succeeded for virtualization service"),
         }
-        Ok(_) => trace!("statslog_rust succeeded for virtualization service"),
-    }
+    });
 }
 
 /// Write the stats of VM boot to statsd
+/// The function creates a separate thread which waits fro statsd to start to push atom
 pub fn write_vm_booted_stats(
     uid: i32,
-    vm_identifier: &String,
+    vm_identifier: &str,
     vm_start_timestamp: Option<SystemTime>,
 ) {
+    let vm_identifier = vm_identifier.to_owned();
     let duration = get_duration(vm_start_timestamp);
-    let vm_booted = vm_booted::VmBooted {
-        uid,
-        vm_identifier,
-        elapsed_time_millis: duration.as_millis() as i64,
-    };
-    match vm_booted.stats_write() {
-        Err(e) => {
-            warn!("statslog_rust failed with error: {}", e);
+    thread::spawn(move || {
+        let vm_booted = vm_booted::VmBooted {
+            uid,
+            vm_identifier: &vm_identifier,
+            elapsed_time_millis: duration.as_millis() as i64,
+        };
+        wait_for_statsd().unwrap_or_else(|e| warn!("failed to wait for statsd with error: {}", e));
+        match vm_booted.stats_write() {
+            Err(e) => {
+                warn!("statslog_rust failed with error: {}", e);
+            }
+            Ok(_) => trace!("statslog_rust succeeded for virtualization service"),
         }
-        Ok(_) => trace!("statslog_rust succeeded for virtualization service"),
-    }
+    });
 }
 
 /// Write the stats of VM exit to statsd
+/// The function creates a separate thread which waits fro statsd to start to push atom
 pub fn write_vm_exited_stats(
     uid: i32,
-    vm_identifier: &String,
+    vm_identifier: &str,
     reason: DeathReason,
     vm_start_timestamp: Option<SystemTime>,
 ) {
+    let vm_identifier = vm_identifier.to_owned();
     let duration = get_duration(vm_start_timestamp);
-    let vm_exited = vm_exited::VmExited {
-        uid,
-        vm_identifier,
-        elapsed_time_millis: duration.as_millis() as i64,
-        death_reason: match reason {
-            DeathReason::INFRASTRUCTURE_ERROR => vm_exited::DeathReason::InfrastructureError,
-            DeathReason::KILLED => vm_exited::DeathReason::Killed,
-            DeathReason::UNKNOWN => vm_exited::DeathReason::Unknown,
-            DeathReason::SHUTDOWN => vm_exited::DeathReason::Shutdown,
-            DeathReason::ERROR => vm_exited::DeathReason::Error,
-            DeathReason::REBOOT => vm_exited::DeathReason::Reboot,
-            DeathReason::CRASH => vm_exited::DeathReason::Crash,
-            DeathReason::PVM_FIRMWARE_PUBLIC_KEY_MISMATCH => {
-                vm_exited::DeathReason::PvmFirmwarePublicKeyMismatch
+    thread::spawn(move || {
+        let vm_exited = vm_exited::VmExited {
+            uid,
+            vm_identifier: &vm_identifier,
+            elapsed_time_millis: duration.as_millis() as i64,
+            death_reason: match reason {
+                DeathReason::INFRASTRUCTURE_ERROR => vm_exited::DeathReason::InfrastructureError,
+                DeathReason::KILLED => vm_exited::DeathReason::Killed,
+                DeathReason::UNKNOWN => vm_exited::DeathReason::Unknown,
+                DeathReason::SHUTDOWN => vm_exited::DeathReason::Shutdown,
+                DeathReason::ERROR => vm_exited::DeathReason::Error,
+                DeathReason::REBOOT => vm_exited::DeathReason::Reboot,
+                DeathReason::CRASH => vm_exited::DeathReason::Crash,
+                DeathReason::PVM_FIRMWARE_PUBLIC_KEY_MISMATCH => {
+                    vm_exited::DeathReason::PvmFirmwarePublicKeyMismatch
+                }
+                DeathReason::PVM_FIRMWARE_INSTANCE_IMAGE_CHANGED => {
+                    vm_exited::DeathReason::PvmFirmwareInstanceImageChanged
+                }
+                DeathReason::BOOTLOADER_PUBLIC_KEY_MISMATCH => {
+                    vm_exited::DeathReason::BootloaderPublicKeyMismatch
+                }
+                DeathReason::BOOTLOADER_INSTANCE_IMAGE_CHANGED => {
+                    vm_exited::DeathReason::BootloaderInstanceImageChanged
+                }
+                DeathReason::MICRODROID_FAILED_TO_CONNECT_TO_VIRTUALIZATION_SERVICE => {
+                    vm_exited::DeathReason::MicrodroidFailedToConnectToVirtualizationService
+                }
+                DeathReason::MICRODROID_PAYLOAD_HAS_CHANGED => {
+                    vm_exited::DeathReason::MicrodroidPayloadHasChanged
+                }
+                DeathReason::MICRODROID_PAYLOAD_VERIFICATION_FAILED => {
+                    vm_exited::DeathReason::MicrodroidPayloadVerificationFailed
+                }
+                DeathReason::MICRODROID_INVALID_PAYLOAD_CONFIG => {
+                    vm_exited::DeathReason::MicrodroidInvalidPayloadConfig
+                }
+                DeathReason::MICRODROID_UNKNOWN_RUNTIME_ERROR => {
+                    vm_exited::DeathReason::MicrodroidUnknownRuntimeError
+                }
+                DeathReason::HANGUP => vm_exited::DeathReason::Hangup,
+                _ => vm_exited::DeathReason::Unknown,
+            },
+        };
+        wait_for_statsd().unwrap_or_else(|e| warn!("failed to wait for statsd with error: {}", e));
+        match vm_exited.stats_write() {
+            Err(e) => {
+                warn!("statslog_rust failed with error: {}", e);
             }
-            DeathReason::PVM_FIRMWARE_INSTANCE_IMAGE_CHANGED => {
-                vm_exited::DeathReason::PvmFirmwareInstanceImageChanged
-            }
-            DeathReason::BOOTLOADER_PUBLIC_KEY_MISMATCH => {
-                vm_exited::DeathReason::BootloaderPublicKeyMismatch
-            }
-            DeathReason::BOOTLOADER_INSTANCE_IMAGE_CHANGED => {
-                vm_exited::DeathReason::BootloaderInstanceImageChanged
-            }
-            DeathReason::MICRODROID_FAILED_TO_CONNECT_TO_VIRTUALIZATION_SERVICE => {
-                vm_exited::DeathReason::MicrodroidFailedToConnectToVirtualizationService
-            }
-            DeathReason::MICRODROID_PAYLOAD_HAS_CHANGED => {
-                vm_exited::DeathReason::MicrodroidPayloadHasChanged
-            }
-            DeathReason::MICRODROID_PAYLOAD_VERIFICATION_FAILED => {
-                vm_exited::DeathReason::MicrodroidPayloadVerificationFailed
-            }
-            DeathReason::MICRODROID_INVALID_PAYLOAD_CONFIG => {
-                vm_exited::DeathReason::MicrodroidInvalidPayloadConfig
-            }
-            DeathReason::MICRODROID_UNKNOWN_RUNTIME_ERROR => {
-                vm_exited::DeathReason::MicrodroidUnknownRuntimeError
-            }
-            DeathReason::HANGUP => vm_exited::DeathReason::Hangup,
-            _ => vm_exited::DeathReason::Unknown,
-        },
-    };
-    match vm_exited.stats_write() {
-        Err(e) => {
-            warn!("statslog_rust failed with error: {}", e);
+            Ok(_) => trace!("statslog_rust succeeded for virtualization service"),
         }
-        Ok(_) => trace!("statslog_rust succeeded for virtualization service"),
-    }
+    });
 }
 
-/// Write the stats of VM cpu status to statsd
-pub fn write_vm_cpu_status_stats(
-    uid: i32,
-    vm_identifier: &String,
-    cpu_status: &VirtualMachineCpuStatus,
-) {
-    let vm_cpu_status_reported = vm_cpu_status_reported::VmCpuStatusReported {
-        uid,
-        vm_identifier,
-        cpu_time_user_millis: cpu_status.cpu_time_user,
-        cpu_time_nice_millis: cpu_status.cpu_time_nice,
-        cpu_time_sys_millis: cpu_status.cpu_time_sys,
-        cpu_time_idle_millis: cpu_status.cpu_time_idle,
-    };
-    match vm_cpu_status_reported.stats_write() {
-        Err(e) => {
-            warn!("statslog_rust failed with error: {}", e);
+fn wait_for_statsd() -> Result<()> {
+    let mut prop = system_properties::PropertyWatcher::new("init.svc.statsd")?;
+    loop {
+        prop.wait()?;
+        match system_properties::read("init.svc.statsd")? {
+            Some(s) => {
+                if s == "running" {
+                    break;
+                }
+            }
+            None => {
+                // This case never really happens because
+                // prop.wait() waits for property to be non-null.
+                break;
+            }
         }
-        Ok(_) => trace!("statslog_rust succeeded for virtualization service"),
     }
-}
-
-/// Write the stats of VM memory status to statsd
-pub fn write_vm_mem_status_stats(
-    uid: i32,
-    vm_identifier: &String,
-    mem_status: &VirtualMachineMemStatus,
-) {
-    let vm_mem_status_reported = vm_mem_status_reported::VmMemStatusReported {
-        uid,
-        vm_identifier,
-        mem_total_kb: mem_status.mem_total,
-        mem_free_kb: mem_status.mem_free,
-        mem_available_kb: mem_status.mem_available,
-        mem_buffer_kb: mem_status.mem_buffer,
-        mem_cached_kb: mem_status.mem_cached,
-    };
-    match vm_mem_status_reported.stats_write() {
-        Err(e) => {
-            warn!("statslog_rust failed with error: {}", e);
-        }
-        Ok(_) => trace!("statslog_rust succeeded for virtualization service"),
-    }
+    Ok(())
 }
diff --git a/virtualizationservice/src/crosvm.rs b/virtualizationservice/src/crosvm.rs
index 6f646b7..1b8061e 100644
--- a/virtualizationservice/src/crosvm.rs
+++ b/virtualizationservice/src/crosvm.rs
@@ -546,6 +546,7 @@
     debug!("Preserving FDs {:?}", preserved_fds);
     command.preserved_fds(preserved_fds);
 
+    command.arg("--params").arg("crashkernel=17M");
     print_crosvm_args(&command);
 
     let result = SharedChild::spawn(&mut command)?;
diff --git a/virtualizationservice/src/main.rs b/virtualizationservice/src/main.rs
index cea2747..714bcfd 100644
--- a/virtualizationservice/src/main.rs
+++ b/virtualizationservice/src/main.rs
@@ -24,8 +24,8 @@
 use crate::aidl::{VirtualizationService, BINDER_SERVICE_IDENTIFIER, TEMPORARY_DIRECTORY};
 use android_logger::{Config, FilterBuilder};
 use android_system_virtualizationservice::aidl::android::system::virtualizationservice::IVirtualizationService::BnVirtualizationService;
+use anyhow::{bail, Context, Error};
 use binder::{register_lazy_service, BinderFeatures, ProcessState};
-use anyhow::Error;
 use log::{info, Level};
 use std::fs::{remove_dir_all, remove_file, read_dir};
 
@@ -44,6 +44,7 @@
             ),
     );
 
+    remove_memlock_rlimit().expect("Failed to remove memlock rlimit");
     clear_temporary_files().expect("Failed to delete old temporary files");
 
     let service = VirtualizationService::init();
@@ -53,6 +54,18 @@
     ProcessState::join_thread_pool();
 }
 
+/// Set this PID's RLIMIT_MEMLOCK to RLIM_INFINITY to allow crosvm (a child process) to mlock()
+/// arbitrary amounts of memory. This is necessary for spawning protected VMs.
+fn remove_memlock_rlimit() -> Result<(), Error> {
+    let lim = libc::rlimit { rlim_cur: libc::RLIM_INFINITY, rlim_max: libc::RLIM_INFINITY };
+    // SAFETY - borrowing the new limit struct only
+    match unsafe { libc::setrlimit(libc::RLIMIT_MEMLOCK, &lim) } {
+        0 => Ok(()),
+        -1 => Err(std::io::Error::last_os_error()).context("setrlimit failed"),
+        n => bail!("Unexpected return value from setrlimit(): {}", n),
+    }
+}
+
 /// Remove any files under `TEMPORARY_DIRECTORY`.
 fn clear_temporary_files() -> Result<(), Error> {
     for dir_entry in read_dir(TEMPORARY_DIRECTORY)? {
diff --git a/virtualizationservice/src/payload.rs b/virtualizationservice/src/payload.rs
index 4190fbb..233e74b 100644
--- a/virtualizationservice/src/payload.rs
+++ b/virtualizationservice/src/payload.rs
@@ -27,7 +27,9 @@
 use microdroid_metadata::{ApexPayload, ApkPayload, Metadata, PayloadConfig, PayloadMetadata};
 use microdroid_payload_config::{ApexConfig, VmPayloadConfig};
 use once_cell::sync::OnceCell;
-use packagemanager_aidl::aidl::android::content::pm::IPackageManagerNative::IPackageManagerNative;
+use packagemanager_aidl::aidl::android::content::pm::{
+    IPackageManagerNative::IPackageManagerNative, StagedApexInfo::StagedApexInfo,
+};
 use regex::Regex;
 use serde::Deserialize;
 use serde_xml_rs::from_reader;
@@ -48,7 +50,7 @@
 const PACKAGE_MANAGER_NATIVE_SERVICE: &str = "package_native";
 
 /// Represents the list of APEXes
-#[derive(Clone, Debug, Deserialize)]
+#[derive(Clone, Debug, Deserialize, Eq, PartialEq)]
 struct ApexInfoList {
     #[serde(rename = "apex-info")]
     list: Vec<ApexInfo>,
@@ -58,6 +60,8 @@
 struct ApexInfo {
     #[serde(rename = "moduleName")]
     name: String,
+    #[serde(rename = "versionCode")]
+    version: u64,
     #[serde(rename = "modulePath")]
     path: PathBuf,
 
@@ -101,6 +105,40 @@
             Ok(apex_info_list)
         })
     }
+
+    // Override apex info with the staged one
+    fn override_staged_apex(&mut self, staged_apex_info: &StagedApexInfo) -> Result<()> {
+        let mut need_to_add: Option<ApexInfo> = None;
+        for apex_info in self.list.iter_mut() {
+            if staged_apex_info.moduleName == apex_info.name {
+                if apex_info.is_active && apex_info.is_factory {
+                    // Copy the entry to the end as factory/non-active after the loop
+                    // to keep the factory version. Typically this step is unncessary,
+                    // but some apexes (like sharedlibs) need to be kept even if it's inactive.
+                    need_to_add.replace(ApexInfo { is_active: false, ..apex_info.clone() });
+                    // And make this one as non-factory. Note that this one is still active
+                    // and overridden right below.
+                    apex_info.is_factory = false;
+                }
+                // Active one is overridden with the staged one.
+                if apex_info.is_active {
+                    apex_info.version = staged_apex_info.versionCode as u64;
+                    apex_info.path = PathBuf::from(&staged_apex_info.diskImagePath);
+                    apex_info.has_classpath_jar = staged_apex_info.hasClassPathJars;
+                    apex_info.last_update_seconds = last_updated(&apex_info.path)?;
+                }
+            }
+        }
+        if let Some(info) = need_to_add {
+            self.list.push(info);
+        }
+        Ok(())
+    }
+}
+
+fn last_updated<P: AsRef<Path>>(path: P) -> Result<u64> {
+    let metadata = metadata(path)?;
+    Ok(metadata.modified()?.duration_since(SystemTime::UNIX_EPOCH)?.as_secs())
 }
 
 impl ApexInfo {
@@ -137,19 +175,11 @@
                     .context("Failed to get service when prefer_staged is set.")?;
             let staged =
                 pm.getStagedApexModuleNames().context("getStagedApexModuleNames failed")?;
-            for apex_info in list.list.iter_mut() {
-                if staged.contains(&apex_info.name) {
-                    if let Some(staged_apex_info) =
-                        pm.getStagedApexInfo(&apex_info.name).context("getStagedApexInfo failed")?
-                    {
-                        apex_info.path = PathBuf::from(staged_apex_info.diskImagePath);
-                        apex_info.has_classpath_jar = staged_apex_info.hasClassPathJars;
-                        let metadata = metadata(&apex_info.path)?;
-                        apex_info.last_update_seconds =
-                            metadata.modified()?.duration_since(SystemTime::UNIX_EPOCH)?.as_secs();
-                        // by definition, staged apex can't be a factory apex.
-                        apex_info.is_factory = false;
-                    }
+            for name in staged {
+                if let Some(staged_apex_info) =
+                    pm.getStagedApexInfo(&name).context("getStagedApexInfo failed")?
+                {
+                    list.override_staged_apex(&staged_apex_info)?;
                 }
             }
         }
@@ -243,8 +273,13 @@
     let apex_list = pm.get_apex_list(vm_payload_config.prefer_staged)?;
 
     // collect APEXes from config
-    let apex_infos =
+    let mut apex_infos =
         collect_apex_infos(&apex_list, &vm_payload_config.apexes, app_config.debugLevel);
+
+    // Pass sorted list of apexes. Sorting key shouldn't use `path` because it will change after
+    // reboot with prefer_staged. `last_update_seconds` is added to distinguish "samegrade"
+    // update.
+    apex_infos.sort_by_key(|info| (&info.name, &info.version, &info.last_update_seconds));
     info!("Microdroid payload APEXes: {:?}", apex_infos.iter().map(|ai| &ai.name));
 
     let metadata_file = make_metadata_file(app_config, &apex_infos, temporary_directory)?;
@@ -422,6 +457,7 @@
 #[cfg(test)]
 mod tests {
     use super::*;
+    use tempfile::NamedTempFile;
 
     #[test]
     fn test_find_apex_names_in_classpath() {
@@ -560,4 +596,90 @@
             ]
         );
     }
+
+    #[test]
+    fn test_prefer_staged_apex_with_factory_active_apex() {
+        let single_apex = ApexInfo {
+            name: "foo".to_string(),
+            version: 1,
+            path: PathBuf::from("foo.apex"),
+            is_factory: true,
+            is_active: true,
+            ..Default::default()
+        };
+        let mut apex_info_list = ApexInfoList { list: vec![single_apex.clone()] };
+
+        let staged = NamedTempFile::new().unwrap();
+        apex_info_list
+            .override_staged_apex(&StagedApexInfo {
+                moduleName: "foo".to_string(),
+                versionCode: 2,
+                diskImagePath: staged.path().to_string_lossy().to_string(),
+                ..Default::default()
+            })
+            .expect("should be ok");
+
+        assert_eq!(
+            apex_info_list,
+            ApexInfoList {
+                list: vec![
+                    ApexInfo {
+                        version: 2,
+                        is_factory: false,
+                        path: staged.path().to_owned(),
+                        last_update_seconds: last_updated(staged.path()).unwrap(),
+                        ..single_apex.clone()
+                    },
+                    ApexInfo { is_active: false, ..single_apex },
+                ],
+            }
+        );
+    }
+
+    #[test]
+    fn test_prefer_staged_apex_with_factory_and_inactive_apex() {
+        let factory_apex = ApexInfo {
+            name: "foo".to_string(),
+            version: 1,
+            path: PathBuf::from("foo.apex"),
+            is_factory: true,
+            ..Default::default()
+        };
+        let active_apex = ApexInfo {
+            name: "foo".to_string(),
+            version: 2,
+            path: PathBuf::from("foo.downloaded.apex"),
+            is_active: true,
+            ..Default::default()
+        };
+        let mut apex_info_list =
+            ApexInfoList { list: vec![factory_apex.clone(), active_apex.clone()] };
+
+        let staged = NamedTempFile::new().unwrap();
+        apex_info_list
+            .override_staged_apex(&StagedApexInfo {
+                moduleName: "foo".to_string(),
+                versionCode: 3,
+                diskImagePath: staged.path().to_string_lossy().to_string(),
+                ..Default::default()
+            })
+            .expect("should be ok");
+
+        assert_eq!(
+            apex_info_list,
+            ApexInfoList {
+                list: vec![
+                    // factory apex isn't touched
+                    factory_apex,
+                    // update active one
+                    ApexInfo {
+                        version: 3,
+                        path: staged.path().to_owned(),
+                        last_update_seconds: last_updated(staged.path()).unwrap(),
+                        ..active_apex
+                    },
+                ],
+            }
+        );
+    }
 }