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/javalib/src/android/system/virtualmachine/VirtualMachine.java b/javalib/src/android/system/virtualmachine/VirtualMachine.java
index e750ae9..214e7e6 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachine.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachine.java
@@ -282,14 +282,7 @@
     @GuardedBy("sCreateLock")
     @NonNull
     private static Map<String, WeakReference<VirtualMachine>> getInstancesMap(Context context) {
-        Map<String, WeakReference<VirtualMachine>> instancesMap;
-        if (sInstances.containsKey(context)) {
-            instancesMap = sInstances.get(context);
-        } else {
-            instancesMap = new HashMap<>();
-            sInstances.put(context, instancesMap);
-        }
-        return instancesMap;
+        return sInstances.computeIfAbsent(context, unused -> new HashMap<>());
     }
 
     @NonNull
diff --git a/microdroid/kernel/README.md b/microdroid/kernel/README.md
index 38eca71..44f828c 100644
--- a/microdroid/kernel/README.md
+++ b/microdroid/kernel/README.md
@@ -39,12 +39,12 @@
 
 For ARM64
 ```bash
-tools/bazel run //common-modules/virtual-device:microdroid_aarch64_config menuconfig
+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
+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/pvmfw/Android.bp b/pvmfw/Android.bp
index b644905..71bac72 100644
--- a/pvmfw/Android.bp
+++ b/pvmfw/Android.bp
@@ -9,6 +9,7 @@
     srcs: ["src/main.rs"],
     edition: "2021",
     rustlibs: [
+        "libbuddy_system_allocator",
         "liblog_rust_nostd",
         "libpvmfw_embedded_key",
         "libvmbase",
diff --git a/pvmfw/src/entry.rs b/pvmfw/src/entry.rs
index 1ea95bc..dc2087d 100644
--- a/pvmfw/src/entry.rs
+++ b/pvmfw/src/entry.rs
@@ -14,6 +14,7 @@
 
 //! Low-level entry and exit points of pvmfw.
 
+use crate::heap;
 use crate::helpers;
 use crate::mmio_guard;
 use core::arch::asm;
@@ -52,6 +53,10 @@
     // - only access MMIO once (and while) it has been mapped and configured
     // - only perform logging once the logger has been initialized
     // - only access non-pvmfw memory once (and while) it has been mapped
+
+    // SAFETY - This function should and will only be called once, here.
+    unsafe { heap::init() };
+
     logger::init(LevelFilter::Info).map_err(|_| RebootReason::InternalError)?;
 
     const FDT_MAX_SIZE: usize = helpers::SIZE_2MB;
diff --git a/pvmfw/src/heap.rs b/pvmfw/src/heap.rs
new file mode 100644
index 0000000..bfa8320
--- /dev/null
+++ b/pvmfw/src/heap.rs
@@ -0,0 +1,26 @@
+// 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.
+
+//! Heap implementation.
+
+use buddy_system_allocator::LockedHeap;
+
+#[global_allocator]
+static HEAP_ALLOCATOR: LockedHeap<32> = LockedHeap::<32>::new();
+
+static mut HEAP: [u8; 65536] = [0; 65536];
+
+pub unsafe fn init() {
+    HEAP_ALLOCATOR.lock().init(HEAP.as_mut_ptr() as usize, HEAP.len());
+}
diff --git a/pvmfw/src/main.rs b/pvmfw/src/main.rs
index 870172d..c0bb263 100644
--- a/pvmfw/src/main.rs
+++ b/pvmfw/src/main.rs
@@ -16,10 +16,12 @@
 
 #![no_main]
 #![no_std]
+#![feature(default_alloc_error_handler)]
 
 mod avb;
 mod entry;
 mod exceptions;
+mod heap;
 mod helpers;
 mod mmio_guard;
 mod smccc;
diff --git a/tests/benchmark/src/native/benchmarkbinary.cpp b/tests/benchmark/src/native/benchmarkbinary.cpp
index 18293f3..6321c25 100644
--- a/tests/benchmark/src/native/benchmarkbinary.cpp
+++ b/tests/benchmark/src/native/benchmarkbinary.cpp
@@ -94,11 +94,12 @@
     Result<double> measure_read_rate(const std::string& filename, int64_t fileSizeBytes,
                                      bool is_rand) {
         const int64_t block_count = fileSizeBytes / kBlockSizeBytes;
-        std::vector<uint64_t> offsets;
+        std::vector<uint64_t> offsets(block_count);
+        for (auto i = 0; i < block_count; ++i) {
+            offsets.push_back(i * kBlockSizeBytes);
+        }
         if (is_rand) {
             std::mt19937 rd{std::random_device{}()};
-            offsets.reserve(block_count);
-            for (auto i = 0; i < block_count; ++i) offsets.push_back(i * kBlockSizeBytes);
             std::shuffle(offsets.begin(), offsets.end(), rd);
         }
         char buf[kBlockSizeBytes];
@@ -109,12 +110,7 @@
             return ErrnoError() << "Read: opening " << filename << " failed";
         }
         for (auto i = 0; i < block_count; ++i) {
-            if (is_rand) {
-                if (lseek(fd.get(), offsets[i], SEEK_SET) == -1) {
-                    return ErrnoError() << "failed to lseek";
-                }
-            }
-            auto bytes = read(fd.get(), buf, kBlockSizeBytes);
+            auto bytes = pread(fd, buf, kBlockSizeBytes, offsets[i]);
             if (bytes == 0) {
                 return Error() << "unexpected end of file";
             } else if (bytes == -1) {
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/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 c9df624..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'");
@@ -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);
 
@@ -569,54 +568,49 @@
     @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");
 
@@ -635,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(
@@ -674,8 +671,6 @@
                         stat.getValue().toString());
             }
         }
-
-        shutdownMicrodroid(getDevice(), cid);
     }
 
     @Test
@@ -717,6 +712,7 @@
     public void setUp() throws Exception {
         testIfDeviceIsCapable(getDevice());
         mMetricPrefix = getMetricPrefix() + "microdroid/";
+        mMicrodroidDevice = null;
 
         prepareVirtualizationTestSetup(getDevice());
 
@@ -728,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 dd01867..cc623a8 100644
--- a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
+++ b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
@@ -28,6 +28,7 @@
 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;
@@ -337,7 +338,10 @@
                     }
                 };
         listener.runToFinish(TAG, vm);
-        assertThat(exception.getNow(null)).isNull();
+        Exception e = exception.getNow(null);
+        if (e != null) {
+            throw e;
+        }
         return vmCdis;
     }
 
@@ -438,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 =
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/src/atom.rs b/virtualizationservice/src/atom.rs
index eabb4cc..20f88e7 100644
--- a/virtualizationservice/src/atom.rs
+++ b/virtualizationservice/src/atom.rs
@@ -26,7 +26,9 @@
 use binder::{ParcelFileDescriptor, ThreadState};
 use log::{trace, warn};
 use microdroid_payload_config::VmPayloadConfig;
+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;
 
@@ -86,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,
@@ -104,105 +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"),
+    });
+}
+
+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(())
 }