diff --git a/.prebuilt_info/prebuilt_info_pvmfw_pvmfw_img.asciipb b/.prebuilt_info/prebuilt_info_pvmfw_pvmfw_img.asciipb
index 7b2cad9..2fa867a 100644
--- a/.prebuilt_info/prebuilt_info_pvmfw_pvmfw_img.asciipb
+++ b/.prebuilt_info/prebuilt_info_pvmfw_pvmfw_img.asciipb
@@ -1,6 +1,6 @@
 drops {
   android_build_drop {
-    build_id: "8964254"
+    build_id: "8966862"
     target: "u-boot_pvmfw"
     source_file: "pvmfw.img"
   }
diff --git a/apkdmverity/Android.bp b/apkdmverity/Android.bp
index 06d4500..97a6174 100644
--- a/apkdmverity/Android.bp
+++ b/apkdmverity/Android.bp
@@ -6,7 +6,7 @@
     name: "apkdmverity.defaults",
     crate_name: "apkdmverity",
     srcs: ["src/main.rs"],
-    edition: "2018",
+    edition: "2021",
     prefer_rlib: true,
     rustlibs: [
         "libanyhow",
diff --git a/authfs/Android.bp b/authfs/Android.bp
index 4c5b70e..7788702 100644
--- a/authfs/Android.bp
+++ b/authfs/Android.bp
@@ -8,7 +8,7 @@
     srcs: [
         "src/main.rs",
     ],
-    edition: "2018",
+    edition: "2021",
     rustlibs: [
         "authfs_aidl_interface-rust",
         "libandroid_logger",
diff --git a/authfs/service/Android.bp b/authfs/service/Android.bp
index 943db35..e9eec1e 100644
--- a/authfs/service/Android.bp
+++ b/authfs/service/Android.bp
@@ -7,7 +7,7 @@
     srcs: [
         "src/main.rs",
     ],
-    edition: "2018",
+    edition: "2021",
     rustlibs: [
         "authfs_aidl_interface-rust",
         "libandroid_logger",
diff --git a/authfs/src/fsverity/metadata/Android.bp b/authfs/src/fsverity/metadata/Android.bp
index c988884..3df7519 100644
--- a/authfs/src/fsverity/metadata/Android.bp
+++ b/authfs/src/fsverity/metadata/Android.bp
@@ -20,6 +20,6 @@
         "libauthfs_fsverity_metadata_bindgen",
         "libopenssl",
     ],
-    edition: "2018",
+    edition: "2021",
     apex_available: ["com.android.virt"],
 }
diff --git a/authfs/tests/Android.bp b/authfs/tests/Android.bp
index 990daf3..72685ad 100644
--- a/authfs/tests/Android.bp
+++ b/authfs/tests/Android.bp
@@ -27,7 +27,7 @@
     name: "open_then_run",
     crate_name: "open_then_run",
     srcs: ["open_then_run.rs"],
-    edition: "2018",
+    edition: "2021",
     rustlibs: [
         "libandroid_logger",
         "libanyhow",
diff --git a/compos/Android.bp b/compos/Android.bp
index ec2d737..6aa9d3d 100644
--- a/compos/Android.bp
+++ b/compos/Android.bp
@@ -4,6 +4,7 @@
 
 rust_defaults {
     name: "compsvc_defaults",
+    edition: "2021",
     srcs: ["src/compsvc_main.rs"],
     rustlibs: [
         "android.system.virtualmachineservice-rust",
diff --git a/compos/benchmark/Android.bp b/compos/benchmark/Android.bp
index 9353771..37af87e 100644
--- a/compos/benchmark/Android.bp
+++ b/compos/benchmark/Android.bp
@@ -11,7 +11,7 @@
     static_libs: [
         "androidx.test.runner",
         "androidx.test.ext.junit",
-        "MicroroidDeviceTestHelper",
+        "MicrodroidDeviceTestHelper",
         "truth-prebuilt",
     ],
     platform_apis: true,
diff --git a/compos/common/Android.bp b/compos/common/Android.bp
index 23a1eb9..7a7042e 100644
--- a/compos/common/Android.bp
+++ b/compos/common/Android.bp
@@ -6,7 +6,7 @@
     name: "libcompos_common",
     crate_name: "compos_common",
     srcs: ["lib.rs"],
-    edition: "2018",
+    edition: "2021",
     rustlibs: [
         "android.system.virtualizationservice-rust",
         "compos_aidl_interface-rust",
diff --git a/compos/composd/Android.bp b/compos/composd/Android.bp
index 937aa1a..07a9be3 100644
--- a/compos/composd/Android.bp
+++ b/compos/composd/Android.bp
@@ -5,7 +5,7 @@
 rust_binary {
     name: "composd",
     srcs: ["src/composd_main.rs"],
-    edition: "2018",
+    edition: "2021",
     prefer_rlib: true,
     rustlibs: [
         "android.system.composd-rust",
diff --git a/compos/composd_cmd/Android.bp b/compos/composd_cmd/Android.bp
index c230e13..61df328 100644
--- a/compos/composd_cmd/Android.bp
+++ b/compos/composd_cmd/Android.bp
@@ -5,7 +5,7 @@
 rust_binary {
     name: "composd_cmd",
     srcs: ["composd_cmd.rs"],
-    edition: "2018",
+    edition: "2021",
     rustlibs: [
         "android.system.composd-rust",
         "libanyhow",
diff --git a/compos/tests/Android.bp b/compos/tests/Android.bp
index 72da97f..41958ca 100644
--- a/compos/tests/Android.bp
+++ b/compos/tests/Android.bp
@@ -10,11 +10,18 @@
         "compatibility-tradefed",
         "compatibility-host-util",
     ],
-    data_native_bins: ["bcc_validator"],
     static_libs: [
         "MicrodroidHostTestHelper",
     ],
     test_suites: [
         "general-tests",
     ],
+    // java_test_host doesn't have data_native_libs but jni_libs can be used to put
+    // native modules under ./lib directory.
+    // This works because host tools have rpath (../lib and ./lib).
+    data_native_bins: ["bcc_validator"],
+    jni_libs: [
+        "libcrypto",
+        "libc++",
+    ],
 }
diff --git a/compos/verify/Android.bp b/compos/verify/Android.bp
index 5c74e4f..f68cc1b 100644
--- a/compos/verify/Android.bp
+++ b/compos/verify/Android.bp
@@ -5,7 +5,7 @@
 rust_binary {
     name: "compos_verify",
     srcs: ["verify.rs"],
-    edition: "2018",
+    edition: "2021",
     rustlibs: [
         "compos_aidl_interface-rust",
         "libandroid_logger",
diff --git a/javalib/Android.bp b/javalib/Android.bp
index 1df9b51..a6c3b80 100644
--- a/javalib/Android.bp
+++ b/javalib/Android.bp
@@ -2,20 +2,34 @@
     default_applicable_licenses: ["Android-Apache-2.0"],
 }
 
+// Defines our permissions
+android_app {
+    name: "android.system.virtualmachine.res",
+    installable: true,
+    apex_available: ["com.android.virt"],
+    sdk_version: "current",
+}
+
 java_sdk_library {
     name: "android.system.virtualmachine",
     installable: true,
     compile_dex: true,
 
+    jarjar_rules: "jarjar-rules.txt",
+
     srcs: ["src/**/*.java"],
     static_libs: [
         "android.system.virtualizationservice-java",
+        // For android.sysprop.HypervisorProperties
+        "PlatformProperties",
     ],
 
     apex_available: ["com.android.virt"],
     permitted_packages: [
         "android.system.virtualmachine",
         "android.system.virtualizationservice",
+        // android.sysprop.*, renamed by jarjar
+        "com.android.system.virtualmachine.sysprop",
     ],
 }
 
@@ -23,10 +37,3 @@
     name: "android-virtualization-framework-sdk",
     api_dirs: ["32"],
 }
-
-android_app {
-    name: "android.system.virtualmachine.res",
-    installable: true,
-    apex_available: ["com.android.virt"],
-    sdk_version: "current",
-}
diff --git a/javalib/jarjar-rules.txt b/javalib/jarjar-rules.txt
new file mode 100644
index 0000000..dd8ad2d
--- /dev/null
+++ b/javalib/jarjar-rules.txt
@@ -0,0 +1,9 @@
+# Rules for the android.system.virtualmachine java_sdk_library.
+
+# This is the root of the API, everything we care about should be
+# reachable from here.
+# (This gets rid of all the android.sysprop classes we don't use.)
+keep android.system.virtualmachine.VirtualMachineManager
+
+# We statically link PlatformProperties, rename to avoid clashes.
+rule android.sysprop.** com.android.system.virtualmachine.sysprop.@1
diff --git a/libs/apexutil/Android.bp b/libs/apexutil/Android.bp
index 5b55e1c..3bdfc5f 100644
--- a/libs/apexutil/Android.bp
+++ b/libs/apexutil/Android.bp
@@ -7,7 +7,7 @@
     crate_name: "apexutil",
     host_supported: true,
     srcs: ["src/lib.rs"],
-    edition: "2018",
+    edition: "2021",
     rustlibs: [
         "liblog_rust",
         "libthiserror",
diff --git a/libs/apkverify/Android.bp b/libs/apkverify/Android.bp
index 2d58bb0..ab9265d 100644
--- a/libs/apkverify/Android.bp
+++ b/libs/apkverify/Android.bp
@@ -7,7 +7,7 @@
     crate_name: "apkverify",
     srcs: ["src/lib.rs"],
     prefer_rlib: true,
-    edition: "2018",
+    edition: "2021",
     rustlibs: [
         "libanyhow",
         "libbyteorder",
@@ -37,7 +37,7 @@
     crate_name: "apkverify_test",
     srcs: ["tests/*_test.rs"],
     prefer_rlib: true,
-    edition: "2018",
+    edition: "2021",
     test_suites: ["general-tests"],
     rustlibs: [
         "libapkverify",
diff --git a/libs/avb_bindgen/Android.bp b/libs/avb_bindgen/Android.bp
index 6dc16e8..80b96a6 100644
--- a/libs/avb_bindgen/Android.bp
+++ b/libs/avb_bindgen/Android.bp
@@ -7,6 +7,7 @@
     host_supported: true,
     wrapper_src: "bindgen/avb.h",
     crate_name: "avb_bindgen",
+    edition: "2021",
     visibility: ["//packages/modules/Virtualization:__subpackages__"],
     source_stem: "bindings",
     bindgen_flags: [
@@ -26,6 +27,7 @@
     name: "libavb_bindgen_test",
     srcs: [":libavb_bindgen"],
     crate_name: "avb_bindgen_test",
+    edition: "2021",
     test_suites: ["general-tests"],
     auto_gen_config: true,
     clippy_lints: "none",
diff --git a/libs/idsig/Android.bp b/libs/idsig/Android.bp
index 25eeae4..9f7d377 100644
--- a/libs/idsig/Android.bp
+++ b/libs/idsig/Android.bp
@@ -6,7 +6,7 @@
     name: "libidsig.defaults",
     crate_name: "idsig",
     srcs: ["src/lib.rs"],
-    edition: "2018",
+    edition: "2021",
     prefer_rlib: true,
     rustlibs: [
         "libanyhow",
diff --git a/libs/nested_virt/Android.bp b/libs/nested_virt/Android.bp
index e364a2d..72393ea 100644
--- a/libs/nested_virt/Android.bp
+++ b/libs/nested_virt/Android.bp
@@ -6,7 +6,7 @@
     name: "libnested_virt",
     crate_name: "nested_virt",
     srcs: ["src/lib.rs"],
-    edition: "2018",
+    edition: "2021",
     rustlibs: [
         "libanyhow",
         "librustutils",
diff --git a/libs/statslog_virtualization/Android.bp b/libs/statslog_virtualization/Android.bp
index 51a51a3..a702ea1 100644
--- a/libs/statslog_virtualization/Android.bp
+++ b/libs/statslog_virtualization/Android.bp
@@ -27,21 +27,28 @@
     ],
 }
 
+rust_defaults {
+    name: "libstatslog_virtualization_rust_defaults",
+    edition: "2021",
+    rustlibs: [
+        "libstatspull_bindgen",
+    ],
+    apex_available: [
+        "com.android.virt",
+    ],
+}
+
 rust_library {
     name: "libstatslog_virtualization_rust_header",
+    defaults: ["libstatslog_virtualization_rust_defaults"],
     crate_name: "statslog_virtualization_rust_header",
     srcs: [
         "statslog_header_wrapper.rs",
         ":statslog_virtualization_header.rs",
     ],
     rustlibs: [
-        "libstatspull_bindgen",
         "libthiserror",
     ],
-    apex_available: [
-        "com.android.virt",
-    ],
-
 }
 
 genrule {
@@ -55,6 +62,7 @@
 
 rust_library {
     name: "libstatslog_virtualization_rust",
+    defaults: ["libstatslog_virtualization_rust_defaults"],
     crate_name: "statslog_virtualization_rust",
     srcs: [
         "statslog_wrapper.rs",
@@ -62,10 +70,5 @@
     ],
     rustlibs: [
         "libstatslog_virtualization_rust_header",
-        "libstatspull_bindgen",
     ],
-    apex_available: [
-        "com.android.virt",
-    ],
-
 }
diff --git a/libs/vbmeta/Android.bp b/libs/vbmeta/Android.bp
index 84dde11..c5078c2 100644
--- a/libs/vbmeta/Android.bp
+++ b/libs/vbmeta/Android.bp
@@ -7,7 +7,7 @@
     crate_name: "vbmeta",
     host_supported: true,
     srcs: ["src/lib.rs"],
-    edition: "2018",
+    edition: "2021",
     rustlibs: [
         "libavb_bindgen",
         "libthiserror",
diff --git a/libs/vmconfig/Android.bp b/libs/vmconfig/Android.bp
index 1aee1ce..fe541d3 100644
--- a/libs/vmconfig/Android.bp
+++ b/libs/vmconfig/Android.bp
@@ -6,7 +6,7 @@
     name: "libvmconfig",
     crate_name: "vmconfig",
     srcs: ["src/lib.rs"],
-    edition: "2018",
+    edition: "2021",
     rustlibs: [
         "android.system.virtualizationservice-rust",
         "libanyhow",
diff --git a/microdroid/payload/config/Android.bp b/microdroid/payload/config/Android.bp
index 827f6e3..7e60cd4 100644
--- a/microdroid/payload/config/Android.bp
+++ b/microdroid/payload/config/Android.bp
@@ -8,7 +8,7 @@
     crate_name: "microdroid_payload_config",
     srcs: ["src/lib.rs"],
     prefer_rlib: true,
-    edition: "2018",
+    edition: "2021",
     rustlibs: [
         "libserde_json",
         "libserde",
diff --git a/microdroid/payload/metadata/Android.bp b/microdroid/payload/metadata/Android.bp
index e4c7692..cd182fc 100644
--- a/microdroid/payload/metadata/Android.bp
+++ b/microdroid/payload/metadata/Android.bp
@@ -8,7 +8,7 @@
     crate_name: "microdroid_metadata",
     srcs: ["src/lib.rs"],
     prefer_rlib: true,
-    edition: "2018",
+    edition: "2021",
     rustlibs: [
         "libanyhow",
         "libmicrodroid_metadata_proto_rust",
diff --git a/microdroid_manager/Android.bp b/microdroid_manager/Android.bp
index da63434..cfb851b 100644
--- a/microdroid_manager/Android.bp
+++ b/microdroid_manager/Android.bp
@@ -6,7 +6,7 @@
     name: "microdroid_manager_defaults",
     crate_name: "microdroid_manager",
     srcs: ["src/main.rs"],
-    edition: "2018",
+    edition: "2021",
     prefer_rlib: true,
     rustlibs: [
         "android.hardware.security.dice-V1-rust",
diff --git a/pvmfw/idmap.S b/pvmfw/idmap.S
index 62555f9..ec3ceaf 100644
--- a/pvmfw/idmap.S
+++ b/pvmfw/idmap.S
@@ -44,9 +44,9 @@
 	.fill		509, 8, 0x0			// 509 GB of remaining VA space
 
 	/* level 2 */
-0:	.fill		511, 8, 0x0
-	.quad		.L_BLOCK_MEM_XIP | 0x7fe00000	// pVM firmware image
+0:	.fill		510, 8, 0x0
+	.quad		.L_BLOCK_MEM_XIP | 0x7fc00000	// pVM firmware image
+	.quad		.L_BLOCK_MEM	 | 0x7fe00000	// Writable memory for stack, heap &c.
 1:	.quad		.L_BLOCK_RO	 | 0x80000000	// DT provided by VMM
 	.quad		.L_BLOCK_RO	 | 0x80200000	// 2 MB of DRAM containing payload image
-	.quad		.L_BLOCK_MEM	 | 0x80400000	// Writable memory for stack, heap &c.
-	.fill		509, 8, 0x0
+	.fill		510, 8, 0x0
diff --git a/pvmfw/image.ld b/pvmfw/image.ld
index aeb5046..18bb3ba 100644
--- a/pvmfw/image.ld
+++ b/pvmfw/image.ld
@@ -16,7 +16,7 @@
 
 MEMORY
 {
-	image		: ORIGIN = 0x7fe00000, LENGTH = 2M
+	image		: ORIGIN = 0x7fc00000, LENGTH = 2M
+	writable_data	: ORIGIN = 0x7fe00000, LENGTH = 2M
 	dtb_region	: ORIGIN = 0x80000000, LENGTH = 2M
-	writable_data	: ORIGIN = 0x80400000, LENGTH = 2M
 }
diff --git a/pvmfw/pvmfw.img b/pvmfw/pvmfw.img
index 7cc8009..dc0d4d7 100644
--- a/pvmfw/pvmfw.img
+++ b/pvmfw/pvmfw.img
Binary files differ
diff --git a/tests/benchmark/Android.bp b/tests/benchmark/Android.bp
index 817df7d..0d1cd91 100644
--- a/tests/benchmark/Android.bp
+++ b/tests/benchmark/Android.bp
@@ -9,7 +9,7 @@
     ],
     srcs: ["src/java/**/*.java"],
     static_libs: [
-        "MicroroidDeviceTestHelper",
+        "MicrodroidDeviceTestHelper",
         "androidx.test.runner",
         "androidx.test.ext.junit",
         "com.android.microdroid.testservice-java",
diff --git a/tests/benchmark/src/java/com/android/microdroid/benchmark/BenchmarkVmListener.java b/tests/benchmark/src/java/com/android/microdroid/benchmark/BenchmarkVmListener.java
new file mode 100644
index 0000000..eb45a71
--- /dev/null
+++ b/tests/benchmark/src/java/com/android/microdroid/benchmark/BenchmarkVmListener.java
@@ -0,0 +1,66 @@
+/*
+ * 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.
+ */
+
+package com.android.microdroid.benchmark;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import android.os.RemoteException;
+import android.system.virtualmachine.VirtualMachine;
+import android.util.Log;
+
+import com.android.microdroid.test.MicrodroidDeviceTestBase.VmEventListener;
+import com.android.microdroid.testservice.IBenchmarkService;
+
+/**
+ * This VM listener is used in {@link MicrodroidBenchmark} tests to facilitate the communication
+ * between the host and VM via {@link IBenchmarkService}.
+ */
+class BenchmarkVmListener extends VmEventListener {
+    private static final String TAG = "BenchmarkVm";
+
+    interface InnerListener {
+        /** This is invoked when both the payload and {@link IBenchmarkService} are ready. */
+        void onPayloadReady(VirtualMachine vm, IBenchmarkService benchmarkService)
+                throws RemoteException;
+    }
+
+    private final InnerListener mListener;
+
+    private BenchmarkVmListener(InnerListener listener) {
+        mListener = listener;
+    }
+
+    @Override
+    public final void onPayloadReady(VirtualMachine vm) {
+        try {
+            IBenchmarkService benchmarkService =
+                    IBenchmarkService.Stub.asInterface(
+                            vm.connectToVsockServer(IBenchmarkService.SERVICE_PORT).get());
+            assertThat(benchmarkService).isNotNull();
+
+            mListener.onPayloadReady(vm, benchmarkService);
+        } catch (Exception e) {
+            Log.e(TAG, "Error inside onPayloadReady():" + e);
+            throw new RuntimeException(e);
+        }
+        forceStop(vm);
+    }
+
+    static BenchmarkVmListener create(InnerListener listener) {
+        return new BenchmarkVmListener(listener);
+    }
+}
diff --git a/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java b/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java
index e19c72a..200b7bf 100644
--- a/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java
+++ b/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java
@@ -24,6 +24,7 @@
 import android.app.Instrumentation;
 import android.os.Bundle;
 import android.os.ParcelFileDescriptor;
+import android.os.RemoteException;
 import android.system.virtualmachine.VirtualMachine;
 import android.system.virtualmachine.VirtualMachineConfig;
 import android.system.virtualmachine.VirtualMachineConfig.DebugLevel;
@@ -193,8 +194,7 @@
             String vmName = "test_vm_io_" + i;
             mInner.forceCreateNewVirtualMachine(vmName, config);
             VirtualMachine vm = mInner.getVirtualMachineManager().get(vmName);
-            VsockVmEventListener listener = new VsockVmEventListener(transferRates, port);
-            listener.runToFinish(TAG, vm);
+            BenchmarkVmListener.create(new VsockListener(transferRates, port)).runToFinish(TAG, vm);
         }
         reportMetrics(transferRates, "vsock/transfer_host_to_vm_", "_mb_per_sec");
     }
@@ -226,8 +226,8 @@
             String vmName = "test_vm_io_" + i;
             mInner.forceCreateNewVirtualMachine(vmName, config);
             VirtualMachine vm = mInner.getVirtualMachineManager().get(vmName);
-            VirtioBlkVmEventListener listener = new VirtioBlkVmEventListener(readRates, isRand);
-            listener.runToFinish(TAG, vm);
+            BenchmarkVmListener.create(new VirtioBlkListener(readRates, isRand))
+                    .runToFinish(TAG, vm);
         }
         reportMetrics(
                 readRates,
@@ -260,14 +260,14 @@
         mInstrumentation.sendStatus(0, bundle);
     }
 
-    private static class VirtioBlkVmEventListener extends VmEventListener {
+    private static class VirtioBlkListener implements BenchmarkVmListener.InnerListener {
         private static final String FILENAME = APEX_ETC_FS + "microdroid_super.img";
 
         private final long mFileSizeBytes;
         private final List<Double> mReadRates;
         private final boolean mIsRand;
 
-        VirtioBlkVmEventListener(List<Double> readRates, boolean isRand) {
+        VirtioBlkListener(List<Double> readRates, boolean isRand) {
             File file = new File(FILENAME);
             try {
                 mFileSizeBytes = Files.size(file.toPath());
@@ -280,32 +280,26 @@
         }
 
         @Override
-        public void onPayloadReady(VirtualMachine vm) {
-            try {
-                IBenchmarkService benchmarkService =
-                        IBenchmarkService.Stub.asInterface(
-                                vm.connectToVsockServer(IBenchmarkService.SERVICE_PORT).get());
-                double elapsedSeconds =
-                        benchmarkService.readFile(FILENAME, mFileSizeBytes, mIsRand);
-                double fileSizeMb = mFileSizeBytes / SIZE_MB;
-                mReadRates.add(fileSizeMb / elapsedSeconds);
-            } catch (Exception e) {
-                throw new RuntimeException(e);
-            }
-            forceStop(vm);
+        public void onPayloadReady(VirtualMachine vm, IBenchmarkService benchmarkService)
+                throws RemoteException {
+            double elapsedSeconds = benchmarkService.readFile(FILENAME, mFileSizeBytes, mIsRand);
+            double fileSizeMb = mFileSizeBytes / SIZE_MB;
+            mReadRates.add(fileSizeMb / elapsedSeconds);
         }
     }
 
     @Test
     public void testMemoryUsage() throws Exception {
         final String vmName = "test_vm_mem_usage";
-        VirtualMachineConfig.Builder builder = mInner.newVmConfigBuilder(
-                "assets/vm_config_io.json");
-        VirtualMachineConfig config = builder.debugLevel(DebugLevel.NONE).memoryMib(256).build();
+        VirtualMachineConfig config =
+                mInner.newVmConfigBuilder("assets/vm_config_io.json")
+                        .debugLevel(DebugLevel.NONE)
+                        .memoryMib(256)
+                        .build();
         mInner.forceCreateNewVirtualMachine(vmName, config);
         VirtualMachine vm = mInner.getVirtualMachineManager().get(vmName);
         MemoryUsageListener listener = new MemoryUsageListener();
-        listener.runToFinish(TAG, vm);
+        BenchmarkVmListener.create(listener).runToFinish(TAG, vm);
 
         double mem_overall = 256.0;
         double mem_total = (double) listener.mMemTotal / 1024.0;
@@ -329,7 +323,7 @@
         mInstrumentation.sendStatus(0, bundle);
     }
 
-    private static class MemoryUsageListener extends VmEventListener {
+    private static class MemoryUsageListener implements BenchmarkVmListener.InnerListener {
         public long mMemTotal;
         public long mMemFree;
         public long mMemAvailable;
@@ -338,55 +332,38 @@
         public long mSlab;
 
         @Override
-        public void onPayloadReady(VirtualMachine vm) {
-            try {
-                IBenchmarkService service =
-                        IBenchmarkService.Stub.asInterface(
-                                vm.connectToVsockServer(IBenchmarkService.SERVICE_PORT).get());
-
-                mMemTotal = service.getMemInfoEntry("MemTotal");
-                mMemFree = service.getMemInfoEntry("MemFree");
-                mMemAvailable = service.getMemInfoEntry("MemAvailable");
-                mBuffers = service.getMemInfoEntry("Buffers");
-                mCached = service.getMemInfoEntry("Cached");
-                mSlab = service.getMemInfoEntry("Slab");
-            } catch (Exception e) {
-                throw new RuntimeException(e);
-            }
-            forceStop(vm);
+        public void onPayloadReady(VirtualMachine vm, IBenchmarkService service)
+                throws RemoteException {
+            mMemTotal = service.getMemInfoEntry("MemTotal");
+            mMemFree = service.getMemInfoEntry("MemFree");
+            mMemAvailable = service.getMemInfoEntry("MemAvailable");
+            mBuffers = service.getMemInfoEntry("Buffers");
+            mCached = service.getMemInfoEntry("Cached");
+            mSlab = service.getMemInfoEntry("Slab");
         }
     }
 
-    private static class VsockVmEventListener extends VmEventListener {
+    private static class VsockListener implements BenchmarkVmListener.InnerListener {
         private static final int NUM_BYTES_TO_TRANSFER = 48 * 1024 * 1024;
 
         private final List<Double> mReadRates;
         private final int mPort;
 
-        VsockVmEventListener(List<Double> readRates, int port) {
+        VsockListener(List<Double> readRates, int port) {
             mReadRates = readRates;
             mPort = port;
         }
 
         @Override
-        public void onPayloadReady(VirtualMachine vm) {
-            try {
-                IBenchmarkService benchmarkService =
-                        IBenchmarkService.Stub.asInterface(
-                                vm.connectToVsockServer(IBenchmarkService.SERVICE_PORT).get());
-                assertThat(benchmarkService).isNotNull();
-                AtomicReference<Double> sendRate = new AtomicReference();
+        public void onPayloadReady(VirtualMachine vm, IBenchmarkService benchmarkService)
+                throws RemoteException {
+            AtomicReference<Double> sendRate = new AtomicReference();
 
-                int serverFd = benchmarkService.initVsockServer(mPort);
-                new Thread(() -> sendRate.set(runVsockClientAndSendData(vm))).start();
-                benchmarkService.runVsockServerAndReceiveData(serverFd, NUM_BYTES_TO_TRANSFER);
+            int serverFd = benchmarkService.initVsockServer(mPort);
+            new Thread(() -> sendRate.set(runVsockClientAndSendData(vm))).start();
+            benchmarkService.runVsockServerAndReceiveData(serverFd, NUM_BYTES_TO_TRANSFER);
 
-                mReadRates.add(sendRate.get());
-            } catch (Exception e) {
-                Log.e(TAG, "Test failed in VM:" + e);
-                throw new RuntimeException(e);
-            }
-            forceStop(vm);
+            mReadRates.add(sendRate.get());
         }
 
         private double runVsockClientAndSendData(VirtualMachine vm) {
diff --git a/tests/benchmark_hostside/java/android/avf/test/AVFHostTestCase.java b/tests/benchmark_hostside/java/android/avf/test/AVFHostTestCase.java
index fc18046..8e9da55 100644
--- a/tests/benchmark_hostside/java/android/avf/test/AVFHostTestCase.java
+++ b/tests/benchmark_hostside/java/android/avf/test/AVFHostTestCase.java
@@ -23,6 +23,8 @@
 import static com.google.common.truth.Truth.assertWithMessage;
 import static com.google.common.truth.TruthJUnit.assume;
 
+import static org.junit.Assume.assumeTrue;
+
 import com.android.microdroid.test.CommandRunner;
 import com.android.microdroid.test.MicrodroidHostTestCaseBase;
 import com.android.tradefed.log.LogUtil.CLog;
@@ -57,7 +59,7 @@
     private static final int BOOT_COMPLETE_TIMEOUT_MS = 10 * 60 * 1000;
     private static final double NANOS_IN_SEC = 1_000_000_000.0;
     private static final int ROUND_COUNT = 5;
-    private static final String METRIC_PREFIX = "avf_perf/compos/";
+    private static final String METRIC_PREFIX = "avf_perf/hostside/";
 
     @Before
     public void setUp() throws Exception {
@@ -66,8 +68,9 @@
 
     @After
     public void tearDown() throws Exception {
-        // Reboot to prevent previous staged session.
-        rebootAndWaitBootCompleted();
+        // Set PKVM enable and reboot to prevent previous staged session.
+        setPKVMStatusWithRebootToBootloader(true);
+        rebootFromBootloaderAndWaitBootCompleted();
 
         CommandRunner android = new CommandRunner(getDevice());
 
@@ -76,6 +79,37 @@
     }
 
     @Test
+    public void testBootEnableAndDisablePKVM() throws Exception {
+
+        testPKVMStatusSwitchSupported();
+
+        double[] bootWithPKVMEnableTime = new double[ROUND_COUNT];
+        double[] bootWithoutPKVMEnableTime = new double[ROUND_COUNT];
+
+        for (int round = 0; round < ROUND_COUNT; ++round) {
+
+            setPKVMStatusWithRebootToBootloader(true);
+            long start = System.nanoTime();
+            rebootFromBootloaderAndWaitBootCompleted();
+            long elapsedWithPKVMEnable = System.nanoTime() - start;
+            double elapsedSec = elapsedWithPKVMEnable / NANOS_IN_SEC;
+            bootWithPKVMEnableTime[round] = elapsedSec;
+            CLog.i("Boot time with PKVM enable took " + elapsedSec + "s");
+
+            setPKVMStatusWithRebootToBootloader(false);
+            start = System.nanoTime();
+            rebootFromBootloaderAndWaitBootCompleted();
+            long elapsedWithoutPKVMEnable = System.nanoTime() - start;
+            elapsedSec = elapsedWithoutPKVMEnable / NANOS_IN_SEC;
+            bootWithoutPKVMEnableTime[round] = elapsedSec;
+            CLog.i("Boot time with PKVM disable took " + elapsedSec + "s");
+        }
+
+        reportMetric("boot_time_with_pkvm_enable", "s", bootWithPKVMEnableTime);
+        reportMetric("boot_time_with_pkvm_disable", "s", bootWithoutPKVMEnableTime);
+    }
+
+    @Test
     public void testBootWithAndWithoutCompOS() throws Exception {
         assume().withMessage("Skip on CF; too slow").that(isCuttlefish()).isFalse();
 
@@ -108,6 +142,18 @@
         reportMetric("boot_time_without_compos", "s", bootWithoutCompOsTime);
     }
 
+    private void testPKVMStatusSwitchSupported() throws Exception {
+        if (!getDevice().isStateBootloaderOrFastbootd()) {
+            getDevice().rebootIntoBootloader();
+        }
+        getDevice().waitForDeviceBootloader();
+
+        CommandResult result;
+        result = getDevice().executeFastbootCommand("oem", "pkvm", "status");
+        rebootFromBootloaderAndWaitBootCompleted();
+        assumeTrue(!result.getStderr().contains("Invalid oem command"));
+    }
+
     private void reportMetric(String name, String unit, double[] values) {
         double sum = 0;
         double min = Double.MAX_VALUE;
@@ -135,6 +181,41 @@
         metrics.addTestMetric(METRIC_PREFIX + name + "_stdev_" + unit, Double.toString(stdev));
     }
 
+    private void setPKVMStatusWithRebootToBootloader(boolean isEnable) throws Exception {
+
+        if (!getDevice().isStateBootloaderOrFastbootd()) {
+            getDevice().rebootIntoBootloader();
+        }
+        getDevice().waitForDeviceBootloader();
+
+        CommandResult result;
+        if (isEnable) {
+            result = getDevice().executeFastbootCommand("oem", "pkvm", "enable");
+        } else {
+            result = getDevice().executeFastbootCommand("oem", "pkvm", "disable");
+        }
+
+        result = getDevice().executeFastbootCommand("oem", "pkvm", "status");
+        CLog.i("Gets PKVM status : " + result);
+
+        String expectedOutput = "";
+
+        if (isEnable) {
+            expectedOutput = "pkvm is enabled";
+        } else {
+            expectedOutput = "pkvm is disabled";
+        }
+        assertWithMessage("Failed to set PKVM status. Reason: " + result)
+            .that(result.toString()).ignoringCase().contains(expectedOutput);
+    }
+
+    private void rebootFromBootloaderAndWaitBootCompleted() throws Exception {
+        getDevice().executeFastbootCommand("reboot");
+        getDevice().waitForDeviceOnline(BOOT_COMPLETE_TIMEOUT_MS);
+        getDevice().waitForBootComplete(BOOT_COMPLETE_TIMEOUT_MS);
+        getDevice().enableAdbRoot();
+    }
+
     private void rebootAndWaitBootCompleted() throws Exception {
         getDevice().nonBlockingReboot();
         getDevice().waitForDeviceOnline(BOOT_COMPLETE_TIMEOUT_MS);
diff --git a/tests/helper/Android.bp b/tests/helper/Android.bp
index 200eb42..f8dcef7 100644
--- a/tests/helper/Android.bp
+++ b/tests/helper/Android.bp
@@ -9,7 +9,7 @@
 }
 
 java_library_static {
-    name: "MicroroidDeviceTestHelper",
+    name: "MicrodroidDeviceTestHelper",
     srcs: ["src/java/com/android/microdroid/**/*.java"],
     static_libs: [
         "androidx.test.runner",
diff --git a/tests/helper/src/java/com/android/microdroid/test/MicrodroidDeviceTestBase.java b/tests/helper/src/java/com/android/microdroid/test/MicrodroidDeviceTestBase.java
index 1a573bb..a80111f 100644
--- a/tests/helper/src/java/com/android/microdroid/test/MicrodroidDeviceTestBase.java
+++ b/tests/helper/src/java/com/android/microdroid/test/MicrodroidDeviceTestBase.java
@@ -118,7 +118,7 @@
         mInner = new Inner(context, protectedVm, VirtualMachineManager.getInstance(context));
     }
 
-    protected abstract static class VmEventListener implements VirtualMachineCallback {
+    public abstract static class VmEventListener implements VirtualMachineCallback {
         private ExecutorService mExecutorService = Executors.newSingleThreadExecutor();
         private OptionalLong mVcpuStartedNanoTime = OptionalLong.empty();
         private OptionalLong mKernelStartedNanoTime = OptionalLong.empty();
diff --git a/tests/testapk/Android.bp b/tests/testapk/Android.bp
index 60912ea..325ebdb 100644
--- a/tests/testapk/Android.bp
+++ b/tests/testapk/Android.bp
@@ -10,7 +10,7 @@
     ],
     srcs: ["src/java/**/*.java"],
     static_libs: [
-        "MicroroidDeviceTestHelper",
+        "MicrodroidDeviceTestHelper",
         "androidx.test.runner",
         "androidx.test.ext.junit",
         "authfs_test_apk_assets",
diff --git a/virtualizationservice/Android.bp b/virtualizationservice/Android.bp
index b31d4a2..b6b7b09 100644
--- a/virtualizationservice/Android.bp
+++ b/virtualizationservice/Android.bp
@@ -6,7 +6,7 @@
     name: "virtualizationservice_defaults",
     crate_name: "virtualizationservice",
     srcs: ["src/main.rs"],
-    edition: "2018",
+    edition: "2021",
     // Only build on targets which crosvm builds on.
     enabled: false,
     target: {
diff --git a/vm/Android.bp b/vm/Android.bp
index 2e914c4..eac640e 100644
--- a/vm/Android.bp
+++ b/vm/Android.bp
@@ -6,7 +6,7 @@
     name: "vm",
     crate_name: "vm",
     srcs: ["src/main.rs"],
-    edition: "2018",
+    edition: "2021",
     prefer_rlib: true,
     rustlibs: [
         "android.system.virtualizationservice-rust",
diff --git a/vmbase/Android.bp b/vmbase/Android.bp
index fb7f6a6..ab87053 100644
--- a/vmbase/Android.bp
+++ b/vmbase/Android.bp
@@ -4,6 +4,7 @@
 
 rust_defaults {
     name: "vmbase_rust_defaults",
+    edition: "2021",
     host_supported: false,
     enabled: false,
     target: {
@@ -55,7 +56,6 @@
     defaults: ["vmbase_rust_defaults"],
     crate_name: "vmbase",
     srcs: ["src/lib.rs"],
-    edition: "2021",
     rustlibs: [
         "liblog_rust_nostd",
         "libpsci",
diff --git a/zipfuse/Android.bp b/zipfuse/Android.bp
index e10fc31..3aba94a 100644
--- a/zipfuse/Android.bp
+++ b/zipfuse/Android.bp
@@ -6,7 +6,7 @@
     name: "zipfuse.defaults",
     crate_name: "zipfuse",
     srcs: ["src/main.rs"],
-    edition: "2018",
+    edition: "2021",
     prefer_rlib: true,
     rustlibs: [
         "libanyhow",