Merge "More tests: connectVsock, no debug"
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachine.java b/javalib/src/android/system/virtualmachine/VirtualMachine.java
index ffcdc51..63b5628 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachine.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachine.java
@@ -49,7 +49,9 @@
 import android.annotation.Nullable;
 import android.annotation.RequiresPermission;
 import android.annotation.SystemApi;
+import android.content.ComponentCallbacks2;
 import android.content.Context;
+import android.content.res.Configuration;
 import android.os.Binder;
 import android.os.IBinder;
 import android.os.ParcelFileDescriptor;
@@ -61,6 +63,7 @@
 import android.system.virtualizationservice.IVirtualMachine;
 import android.system.virtualizationservice.IVirtualMachineCallback;
 import android.system.virtualizationservice.IVirtualizationService;
+import android.system.virtualizationservice.MemoryTrimLevel;
 import android.system.virtualizationservice.PartitionType;
 import android.system.virtualizationservice.VirtualMachineAppConfig;
 import android.system.virtualizationservice.VirtualMachineState;
@@ -193,6 +196,56 @@
      */
     @NonNull private final List<ExtraApkSpec> mExtraApks;
 
+    private class MemoryManagementCallbacks implements ComponentCallbacks2 {
+        @Override
+        public void onConfigurationChanged(@NonNull Configuration newConfig) {}
+
+        @Override
+        public void onLowMemory() {}
+
+        @Override
+        public void onTrimMemory(int level) {
+            @MemoryTrimLevel int vmTrimLevel;
+
+            switch (level) {
+                case ComponentCallbacks2.TRIM_MEMORY_RUNNING_CRITICAL:
+                    vmTrimLevel = MemoryTrimLevel.TRIM_MEMORY_RUNNING_CRITICAL;
+                    break;
+                case ComponentCallbacks2.TRIM_MEMORY_RUNNING_LOW:
+                    vmTrimLevel = MemoryTrimLevel.TRIM_MEMORY_RUNNING_LOW;
+                    break;
+                case ComponentCallbacks2.TRIM_MEMORY_RUNNING_MODERATE:
+                    vmTrimLevel = MemoryTrimLevel.TRIM_MEMORY_RUNNING_MODERATE;
+                    break;
+                case ComponentCallbacks2.TRIM_MEMORY_BACKGROUND:
+                case ComponentCallbacks2.TRIM_MEMORY_MODERATE:
+                case ComponentCallbacks2.TRIM_MEMORY_COMPLETE:
+                    /* Release as much memory as we can. The app is on the LMKD LRU kill list. */
+                    vmTrimLevel = MemoryTrimLevel.TRIM_MEMORY_RUNNING_CRITICAL;
+                    break;
+                default:
+                    /* Treat unrecognised messages as generic low-memory warnings. */
+                    vmTrimLevel = MemoryTrimLevel.TRIM_MEMORY_RUNNING_LOW;
+                    break;
+            }
+
+            synchronized (mLock) {
+                try {
+                    if (mVirtualMachine != null) {
+                        mVirtualMachine.onTrimMemory(vmTrimLevel);
+                    }
+                } catch (Exception e) {
+                    /* Caller doesn't want our exceptions. Log them instead. */
+                    Log.w(TAG, "TrimMemory failed: ", e);
+                }
+            }
+        }
+    }
+
+    @NonNull private final MemoryManagementCallbacks mMemoryManagementCallbacks;
+
+    @NonNull private final Context mContext;
+
     // A note on lock ordering:
     // You can take mLock while holding VirtualMachineManager.sCreateLock, but not vice versa.
     // We never take any other lock while holding mCallbackLock; therefore you can
@@ -268,6 +321,8 @@
         mInstanceFilePath = new File(thisVmDir, INSTANCE_IMAGE_FILE);
         mIdsigFilePath = new File(thisVmDir, IDSIG_FILE);
         mExtraApks = setupExtraApks(context, config, thisVmDir);
+        mMemoryManagementCallbacks = new MemoryManagementCallbacks();
+        mContext = context;
     }
 
     /**
@@ -696,6 +751,7 @@
                                 executeCallback((cb) -> cb.onRamdump(VirtualMachine.this, ramdump));
                             }
                         });
+                mContext.registerComponentCallbacks(mMemoryManagementCallbacks);
                 service.asBinder().linkToDeath(deathRecipient, 0);
                 mVirtualMachine.start();
             } catch (IOException | IllegalStateException | ServiceSpecificException e) {
@@ -772,6 +828,7 @@
             }
             try {
                 mVirtualMachine.stop();
+                mContext.unregisterComponentCallbacks(mMemoryManagementCallbacks);
                 mVirtualMachine = null;
             } catch (RemoteException e) {
                 throw e.rethrowAsRuntimeException();
@@ -797,6 +854,7 @@
             try {
                 if (stateToStatus(mVirtualMachine.getState()) == STATUS_RUNNING) {
                     mVirtualMachine.stop();
+                    mContext.unregisterComponentCallbacks(mMemoryManagementCallbacks);
                     mVirtualMachine = null;
                 }
             } catch (RemoteException e) {
diff --git a/microdroid/init.rc b/microdroid/init.rc
index 7d04557..310cf2b 100644
--- a/microdroid/init.rc
+++ b/microdroid/init.rc
@@ -45,7 +45,7 @@
 
     setprop ro.debuggable ${ro.boot.microdroid.debuggable:-0}
 
-on property:dev.bootcomplete=1
+on property:microdroid_manager.init_done=1
     # Stop ueventd to save memory
     stop ueventd
 
diff --git a/microdroid_manager/aidl/android/system/virtualization/payload/IVmPayloadService.aidl b/microdroid_manager/aidl/android/system/virtualization/payload/IVmPayloadService.aidl
index f8e7d34..3859785 100644
--- a/microdroid_manager/aidl/android/system/virtualization/payload/IVmPayloadService.aidl
+++ b/microdroid_manager/aidl/android/system/virtualization/payload/IVmPayloadService.aidl
@@ -27,6 +27,12 @@
     /** Path to the APK contents path. */
     const String VM_APK_CONTENTS_PATH = "/mnt/apk";
 
+    /**
+     * Path to the encrypted storage. Note the path will not exist if encrypted storage
+     * is not enabled.
+     */
+    const String ENCRYPTEDSTORE_MOUNTPOINT = "/mnt/encryptedstore";
+
     /** Notifies that the payload is ready to serve. */
     void notifyPayloadReady();
 
diff --git a/microdroid_manager/src/main.rs b/microdroid_manager/src/main.rs
index 4f94bb4..0e45461 100644
--- a/microdroid_manager/src/main.rs
+++ b/microdroid_manager/src/main.rs
@@ -29,6 +29,7 @@
 use android_system_virtualization_payload::aidl::android::system::virtualization::payload::IVmPayloadService::{
     VM_APK_CONTENTS_PATH,
     VM_PAYLOAD_SERVICE_SOCKET_NAME,
+    ENCRYPTEDSTORE_MOUNTPOINT,
 };
 use anyhow::{anyhow, bail, ensure, Context, Error, Result};
 use apkverify::{get_public_key_der, verify, V4Signature};
@@ -86,7 +87,6 @@
 const ENCRYPTEDSTORE_BIN: &str = "/system/bin/encryptedstore";
 const ENCRYPTEDSTORE_KEY_IDENTIFIER: &str = "encryptedstore_key";
 const ENCRYPTEDSTORE_KEYSIZE: u32 = 32;
-const ENCRYPTEDSTORE_MOUNTPOINT: &str = "/mnt/encryptedstore";
 
 #[derive(thiserror::Error, Debug)]
 enum MicrodroidError {
@@ -433,12 +433,17 @@
 
     register_vm_payload_service(allow_restricted_apis, service.clone(), dice_context)?;
 
+    // Wait for encryptedstore to finish mounting the storage (if enabled) before setting
+    // microdroid_manager.init_done. Reason is init stops uneventd after that.
+    // Encryptedstore, however requires ueventd
     if let Some(mut child) = encryptedstore_child {
         let exitcode = child.wait().context("Wait for encryptedstore child")?;
         ensure!(exitcode.success(), "Unable to prepare encrypted storage. Exitcode={}", exitcode);
     }
 
     wait_for_property_true("dev.bootcomplete").context("failed waiting for dev.bootcomplete")?;
+    system_properties::write("microdroid_manager.init_done", "1")
+        .context("set microdroid_manager.init_done")?;
     info!("boot completed, time to run payload");
     exec_task(task, service).context("Failed to run payload")
 }
diff --git a/tests/aidl/com/android/microdroid/testservice/IBenchmarkService.aidl b/tests/aidl/com/android/microdroid/testservice/IBenchmarkService.aidl
index 16e4893..c8c8660 100644
--- a/tests/aidl/com/android/microdroid/testservice/IBenchmarkService.aidl
+++ b/tests/aidl/com/android/microdroid/testservice/IBenchmarkService.aidl
@@ -30,6 +30,9 @@
     /** Returns an entry from /proc/meminfo. */
     long getMemInfoEntry(String name);
 
+    /** Allocates anonymous memory and returns the raw pointer. */
+    long allocAnonMemory(long mb);
+
     /**
      * Initializes the vsock server on VM.
      * @return the server socket file descriptor.
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 14a0e39..3c3faf2 100644
--- a/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java
+++ b/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java
@@ -27,6 +27,7 @@
 import android.app.Instrumentation;
 import android.os.Bundle;
 import android.os.ParcelFileDescriptor;
+import android.os.Process;
 import android.os.RemoteException;
 import android.system.virtualmachine.VirtualMachine;
 import android.system.virtualmachine.VirtualMachineConfig;
@@ -279,6 +280,58 @@
         return runInShell(TAG, mInstrumentation.getUiAutomation(), command);
     }
 
+    private static class CrosvmStats {
+        public final long mHostRss;
+        public final long mHostPss;
+        public final long mGuestRss;
+        public final long mGuestPss;
+
+        CrosvmStats(Function<String, String> shellExecutor) {
+            try {
+                List<Integer> crosvmPids =
+                        ProcessUtil.getProcessMap(shellExecutor).entrySet().stream()
+                                .filter(e -> e.getValue().contains("crosvm"))
+                                .map(e -> e.getKey())
+                                .collect(java.util.stream.Collectors.toList());
+                if (crosvmPids.size() != 1) {
+                    throw new IllegalStateException(
+                            "expected to find exactly one crosvm processes, found "
+                                    + crosvmPids.size());
+                }
+
+                long hostRss = 0;
+                long hostPss = 0;
+                long guestRss = 0;
+                long guestPss = 0;
+                boolean hasGuestMaps = false;
+                for (ProcessUtil.SMapEntry entry :
+                        ProcessUtil.getProcessSmaps(crosvmPids.get(0), shellExecutor)) {
+                    long rss = entry.metrics.get("Rss");
+                    long pss = entry.metrics.get("Pss");
+                    if (entry.name.contains("crosvm_guest")) {
+                        guestRss += rss;
+                        guestPss += pss;
+                        hasGuestMaps = true;
+                    } else {
+                        hostRss += rss;
+                        hostPss += pss;
+                    }
+                }
+                if (!hasGuestMaps) {
+                    throw new IllegalStateException(
+                            "found no crosvm_guest smap entry in crosvm process");
+                }
+                mHostRss = hostRss;
+                mHostPss = hostPss;
+                mGuestRss = guestRss;
+                mGuestPss = guestPss;
+            } catch (Exception e) {
+                Log.e(TAG, "Error inside onPayloadReady():" + e);
+                throw new RuntimeException(e);
+            }
+        }
+    }
+
     @Test
     public void testMemoryUsage() throws Exception {
         final String vmName = "test_vm_mem_usage";
@@ -299,10 +352,10 @@
         double mem_buffers = (double) listener.mBuffers / 1024.0;
         double mem_cached = (double) listener.mCached / 1024.0;
         double mem_slab = (double) listener.mSlab / 1024.0;
-        double mem_crosvm_host_rss = (double) listener.mCrosvmHostRss / 1024.0;
-        double mem_crosvm_host_pss = (double) listener.mCrosvmHostPss / 1024.0;
-        double mem_crosvm_guest_rss = (double) listener.mCrosvmGuestRss / 1024.0;
-        double mem_crosvm_guest_pss = (double) listener.mCrosvmGuestPss / 1024.0;
+        double mem_crosvm_host_rss = (double) listener.mCrosvm.mHostRss / 1024.0;
+        double mem_crosvm_host_pss = (double) listener.mCrosvm.mHostPss / 1024.0;
+        double mem_crosvm_guest_rss = (double) listener.mCrosvm.mGuestRss / 1024.0;
+        double mem_crosvm_guest_pss = (double) listener.mCrosvm.mGuestPss / 1024.0;
 
         double mem_kernel = mem_overall - mem_total;
         double mem_used = mem_total - mem_free - mem_buffers - mem_cached - mem_slab;
@@ -327,7 +380,7 @@
             mShellExecutor = shellExecutor;
         }
 
-        public Function<String, String> mShellExecutor;
+        public final Function<String, String> mShellExecutor;
 
         public long mMemTotal;
         public long mMemFree;
@@ -336,10 +389,7 @@
         public long mCached;
         public long mSlab;
 
-        public long mCrosvmHostRss;
-        public long mCrosvmHostPss;
-        public long mCrosvmGuestRss;
-        public long mCrosvmGuestPss;
+        public CrosvmStats mCrosvm;
 
         @Override
         public void onPayloadReady(VirtualMachine vm, IBenchmarkService service)
@@ -350,39 +400,80 @@
             mBuffers = service.getMemInfoEntry("Buffers");
             mCached = service.getMemInfoEntry("Cached");
             mSlab = service.getMemInfoEntry("Slab");
+            mCrosvm = new CrosvmStats(mShellExecutor);
+        }
+    }
 
+    @Test
+    public void testMemoryReclaim() throws Exception {
+        final String vmName = "test_vm_mem_reclaim";
+        VirtualMachineConfig config =
+                newVmConfigBuilder()
+                        .setPayloadConfigPath("assets/vm_config_io.json")
+                        .setDebugLevel(DEBUG_LEVEL_NONE)
+                        .setMemoryMib(256)
+                        .build();
+        VirtualMachine vm = forceCreateNewVirtualMachine(vmName, config);
+        MemoryReclaimListener listener = new MemoryReclaimListener(this::executeCommand);
+        BenchmarkVmListener.create(listener).runToFinish(TAG, vm);
+
+        double mem_pre_crosvm_host_rss = (double) listener.mPreCrosvm.mHostRss / 1024.0;
+        double mem_pre_crosvm_host_pss = (double) listener.mPreCrosvm.mHostPss / 1024.0;
+        double mem_pre_crosvm_guest_rss = (double) listener.mPreCrosvm.mGuestRss / 1024.0;
+        double mem_pre_crosvm_guest_pss = (double) listener.mPreCrosvm.mGuestPss / 1024.0;
+        double mem_post_crosvm_host_rss = (double) listener.mPostCrosvm.mHostRss / 1024.0;
+        double mem_post_crosvm_host_pss = (double) listener.mPostCrosvm.mHostPss / 1024.0;
+        double mem_post_crosvm_guest_rss = (double) listener.mPostCrosvm.mGuestRss / 1024.0;
+        double mem_post_crosvm_guest_pss = (double) listener.mPostCrosvm.mGuestPss / 1024.0;
+
+        Bundle bundle = new Bundle();
+        bundle.putDouble(
+                METRIC_NAME_PREFIX + "mem_pre_crosvm_host_rss_MB", mem_pre_crosvm_host_rss);
+        bundle.putDouble(
+                METRIC_NAME_PREFIX + "mem_pre_crosvm_host_pss_MB", mem_pre_crosvm_host_pss);
+        bundle.putDouble(
+                METRIC_NAME_PREFIX + "mem_pre_crosvm_guest_rss_MB", mem_pre_crosvm_guest_rss);
+        bundle.putDouble(
+                METRIC_NAME_PREFIX + "mem_pre_crosvm_guest_pss_MB", mem_pre_crosvm_guest_pss);
+        bundle.putDouble(
+                METRIC_NAME_PREFIX + "mem_post_crosvm_host_rss_MB", mem_post_crosvm_host_rss);
+        bundle.putDouble(
+                METRIC_NAME_PREFIX + "mem_post_crosvm_host_pss_MB", mem_post_crosvm_host_pss);
+        bundle.putDouble(
+                METRIC_NAME_PREFIX + "mem_post_crosvm_guest_rss_MB", mem_post_crosvm_guest_rss);
+        bundle.putDouble(
+                METRIC_NAME_PREFIX + "mem_post_crosvm_guest_pss_MB", mem_post_crosvm_guest_pss);
+        mInstrumentation.sendStatus(0, bundle);
+    }
+
+    private static class MemoryReclaimListener implements BenchmarkVmListener.InnerListener {
+        MemoryReclaimListener(Function<String, String> shellExecutor) {
+            mShellExecutor = shellExecutor;
+        }
+
+        public final Function<String, String> mShellExecutor;
+
+        public CrosvmStats mPreCrosvm;
+        public CrosvmStats mPostCrosvm;
+
+        @Override
+        @SuppressWarnings("ReturnValueIgnored")
+        public void onPayloadReady(VirtualMachine vm, IBenchmarkService service)
+                throws RemoteException {
+            // Allocate 256MB of anonymous memory. This will fill all guest
+            // memory and cause swapping to start.
+            service.allocAnonMemory(256);
+            mPreCrosvm = new CrosvmStats(mShellExecutor);
+            // Send a memory trim hint to cause memory reclaim.
+            mShellExecutor.apply("am send-trim-memory " + Process.myPid() + " RUNNING_CRITICAL");
+            // Give time for the memory reclaim to do its work.
             try {
-                List<Integer> crosvmPids =
-                        ProcessUtil.getProcessMap(mShellExecutor).entrySet().stream()
-                                .filter(e -> e.getValue().contains("crosvm"))
-                                .map(e -> e.getKey())
-                                .collect(java.util.stream.Collectors.toList());
-                if (crosvmPids.size() != 1) {
-                    throw new RuntimeException(
-                            "expected to find exactly one crosvm processes, found "
-                                    + crosvmPids.size());
-                }
-
-                mCrosvmHostRss = 0;
-                mCrosvmHostPss = 0;
-                mCrosvmGuestRss = 0;
-                mCrosvmGuestPss = 0;
-                for (ProcessUtil.SMapEntry entry :
-                        ProcessUtil.getProcessSmaps(crosvmPids.get(0), mShellExecutor)) {
-                    long rss = entry.metrics.get("Rss");
-                    long pss = entry.metrics.get("Pss");
-                    if (entry.name.contains("crosvm_guest")) {
-                        mCrosvmGuestRss += rss;
-                        mCrosvmGuestPss += pss;
-                    } else {
-                        mCrosvmHostRss += rss;
-                        mCrosvmHostPss += pss;
-                    }
-                }
-            } catch (Exception e) {
-                Log.e(TAG, "Error inside onPayloadReady():" + e);
-                throw new RuntimeException(e);
+                Thread.sleep(isCuttlefish() ? 10000 : 5000);
+            } catch (InterruptedException e) {
+                Log.e(TAG, "Interrupted sleep:" + e);
+                Thread.currentThread().interrupt();
             }
+            mPostCrosvm = new CrosvmStats(mShellExecutor);
         }
     }
 
diff --git a/tests/benchmark/src/native/benchmarkbinary.cpp b/tests/benchmark/src/native/benchmarkbinary.cpp
index 70ec7db..5c172c0 100644
--- a/tests/benchmark/src/native/benchmarkbinary.cpp
+++ b/tests/benchmark/src/native/benchmarkbinary.cpp
@@ -77,6 +77,11 @@
         return ndk::ScopedAStatus::ok();
     }
 
+    ndk::ScopedAStatus allocAnonMemory(long mb, long* out) override {
+        *out = (long)alloc_anon_memory(mb);
+        return ndk::ScopedAStatus::ok();
+    }
+
     ndk::ScopedAStatus initVsockServer(int32_t port, int32_t* out) override {
         auto res = io_vsock::init_vsock_server(port);
         if (res.ok()) {
@@ -131,6 +136,17 @@
         return {file_size_mb / elapsed_seconds};
     }
 
+    void* alloc_anon_memory(long mb) {
+        long bytes = mb << 20;
+        void* p = malloc(bytes);
+        /*
+         * Heap memory is demand allocated. Dirty all pages to ensure
+         * all are allocated.
+         */
+        memset(p, 0x55, bytes);
+        return p;
+    }
+
     Result<size_t> read_meminfo_entry(const std::string& stat) {
         std::ifstream fs("/proc/meminfo");
         if (!fs.is_open()) {
diff --git a/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java b/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
index 11b3e84..231fc7b 100644
--- a/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
+++ b/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
@@ -407,20 +407,26 @@
     }
 
     @Test
-    @Ignore("b/245081929")
     @CddTest(requirements = {"9.17/C-2-1", "9.17/C-2-2", "9.17/C-2-6"})
-    public void testBootFailsWhenProtectedVmStartsWithImagesSignedWithDifferentKey()
-            throws Exception {
+    public void protectedVmWithImageSignedWithDifferentKeyRunsPvmfw() throws Exception {
+        // Arrange
         boolean protectedVm = true;
         assumeTrue(
                 "Skip if protected VMs are not supported",
                 getAndroidDevice().supportsMicrodroid(protectedVm));
-
         File key = findTestFile("test.com.android.virt.pem");
-        Map<String, File> keyOverrides = Map.of();
-        VmInfo vmInfo = runMicrodroidWithResignedImages(key, keyOverrides, protectedVm);
+
+        // Act
+        VmInfo vmInfo =
+                runMicrodroidWithResignedImages(key, /*keyOverrides=*/ Map.of(), protectedVm);
+
+        // Asserts
         vmInfo.mProcess.waitFor(5L, TimeUnit.SECONDS);
-        assertThat(getDevice().pullFileContents(CONSOLE_PATH), containsString("pvmfw boot failed"));
+        String consoleLog = getDevice().pullFileContents(CONSOLE_PATH);
+        assertWithMessage("pvmfw should start").that(consoleLog).contains("pVM firmware");
+        // TODO(b/256148034): Asserts that pvmfw run fails when this verification is implemented.
+        // Also rename the test.
+        vmInfo.mProcess.destroy();
     }
 
     // TODO(b/245277660): Resigning the system/vendor image changes the vbmeta hash.
diff --git a/virtualizationservice/Android.bp b/virtualizationservice/Android.bp
index b767013..da56f76 100644
--- a/virtualizationservice/Android.bp
+++ b/virtualizationservice/Android.bp
@@ -51,6 +51,7 @@
         "libshared_child",
         "libstatslog_virtualization_rust",
         "libtombstoned_client_rust",
+        "libvm_control",
         "libvmconfig",
         "libzip",
         "libvsock",
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualMachine.aidl b/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualMachine.aidl
index d9d9a61..d76b586 100644
--- a/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualMachine.aidl
+++ b/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualMachine.aidl
@@ -16,6 +16,7 @@
 package android.system.virtualizationservice;
 
 import android.system.virtualizationservice.IVirtualMachineCallback;
+import android.system.virtualizationservice.MemoryTrimLevel;
 import android.system.virtualizationservice.VirtualMachineState;
 
 interface IVirtualMachine {
@@ -41,6 +42,9 @@
      */
     void stop();
 
+    /** Communicate app low-memory notifications to the VM. */
+    void onTrimMemory(MemoryTrimLevel level);
+
     /** Open a vsock connection to the CID of the VM on the given port. */
     ParcelFileDescriptor connectVsock(int port);
 }
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice/MemoryTrimLevel.aidl b/virtualizationservice/aidl/android/system/virtualizationservice/MemoryTrimLevel.aidl
new file mode 100644
index 0000000..9ed9e99
--- /dev/null
+++ b/virtualizationservice/aidl/android/system/virtualizationservice/MemoryTrimLevel.aidl
@@ -0,0 +1,27 @@
+/*
+ * Copyright 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package android.system.virtualizationservice;
+
+/**
+ * Memory trim levels propagated from the app to the VM.
+ */
+@Backing(type="int")
+enum MemoryTrimLevel {
+    /* Same meaning as in ComponentCallbacks2 */
+    TRIM_MEMORY_RUNNING_CRITICAL = 0,
+    TRIM_MEMORY_RUNNING_LOW = 1,
+    TRIM_MEMORY_RUNNING_MODERATE = 2,
+}
diff --git a/virtualizationservice/src/aidl.rs b/virtualizationservice/src/aidl.rs
index 040c0d8..7d24a32 100644
--- a/virtualizationservice/src/aidl.rs
+++ b/virtualizationservice/src/aidl.rs
@@ -27,6 +27,7 @@
     IVirtualMachine::{BnVirtualMachine, IVirtualMachine},
     IVirtualMachineCallback::IVirtualMachineCallback,
     IVirtualizationService::IVirtualizationService,
+    MemoryTrimLevel::MemoryTrimLevel,
     Partition::Partition,
     PartitionType::PartitionType,
     VirtualMachineAppConfig::{Payload::Payload, VirtualMachineAppConfig},
@@ -961,6 +962,13 @@
         })
     }
 
+    fn onTrimMemory(&self, level: MemoryTrimLevel) -> binder::Result<()> {
+        self.instance.trim_memory(level).map_err(|e| {
+            error!("Error trimming VM with CID {}: {:?}", self.instance.cid, e);
+            Status::new_service_specific_error_str(-1, Some(e.to_string()))
+        })
+    }
+
     fn connectVsock(&self, port: i32) -> binder::Result<ParcelFileDescriptor> {
         if !matches!(&*self.instance.vm_state.lock().unwrap(), VmState::Running { .. }) {
             return Err(Status::new_service_specific_error_str(-1, Some("VM is not running")));
diff --git a/virtualizationservice/src/crosvm.rs b/virtualizationservice/src/crosvm.rs
index 13e5c70..0fdc293 100644
--- a/virtualizationservice/src/crosvm.rs
+++ b/virtualizationservice/src/crosvm.rs
@@ -22,12 +22,13 @@
 use libc::{sysconf, _SC_CLK_TCK};
 use log::{debug, error, info};
 use semver::{Version, VersionReq};
-use nix::{fcntl::OFlag, unistd::pipe2};
+use nix::{fcntl::OFlag, unistd::pipe2, unistd::Uid, unistd::User};
 use regex::{Captures, Regex};
 use rustutils::system_properties;
 use shared_child::SharedChild;
 use std::borrow::Cow;
 use std::cmp::max;
+use std::fmt;
 use std::fs::{read_to_string, remove_dir_all, File};
 use std::io::{self, Read};
 use std::mem;
@@ -38,7 +39,10 @@
 use std::sync::{Arc, Condvar, Mutex};
 use std::time::{Duration, SystemTime};
 use std::thread;
-use android_system_virtualizationservice::aidl::android::system::virtualizationservice::DeathReason::DeathReason;
+use android_system_virtualizationservice::aidl::android::system::virtualizationservice::{
+    DeathReason::DeathReason,
+    MemoryTrimLevel::MemoryTrimLevel,
+};
 use android_system_virtualizationservice_internal::aidl::android::system::virtualizationservice_internal::IGlobalVmContext::IGlobalVmContext;
 use binder::Strong;
 use android_system_virtualmachineservice::aidl::android::system::virtualmachineservice::IVirtualMachineService::IVirtualMachineService;
@@ -47,6 +51,7 @@
 
 /// external/crosvm
 use base::UnixSeqpacketListener;
+use vm_control::{BalloonControlCommand, VmRequest, VmResponse};
 
 const CROSVM_PATH: &str = "/apex/com.android.virt/bin/crosvm";
 
@@ -171,7 +176,7 @@
 
             // If this fails and returns an error, `self` will be left in the `Failed` state.
             let child =
-                Arc::new(run_vm(config, &instance.temporary_directory, failure_pipe_write)?);
+                Arc::new(run_vm(config, &instance.crosvm_control_socket_path, failure_pipe_write)?);
 
             let instance_monitor_status = instance.clone();
             let child_monitor_status = child.clone();
@@ -228,6 +233,8 @@
     vm_context: VmContext,
     /// The CID assigned to the VM for vsock communication.
     pub cid: Cid,
+    /// Path to crosvm control socket
+    crosvm_control_socket_path: PathBuf,
     /// The name of the VM.
     pub name: String,
     /// Whether the VM is a protected VM.
@@ -249,6 +256,19 @@
     payload_state: Mutex<PayloadState>,
     /// Represents the condition that payload_state was updated
     payload_state_updated: Condvar,
+    /// The human readable name of requester_uid
+    requester_uid_name: String,
+}
+
+impl fmt::Display for VmInstance {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let adj = if self.protected { "Protected" } else { "Non-protected" };
+        write!(
+            f,
+            "{} virtual machine \"{}\" (owner: {}, cid: {})",
+            adj, self.name, self.requester_uid_name, self.cid
+        )
+    }
 }
 
 impl VmInstance {
@@ -264,10 +284,15 @@
         let cid = config.cid;
         let name = config.name.clone();
         let protected = config.protected;
-        Ok(VmInstance {
+        let requester_uid_name = User::from_uid(Uid::from_raw(requester_uid))
+            .ok()
+            .flatten()
+            .map_or_else(|| format!("{}", requester_uid), |u| u.name);
+        let instance = VmInstance {
             vm_state: Mutex::new(VmState::NotStarted { config }),
             vm_context,
             cid,
+            crosvm_control_socket_path: temporary_directory.join("crosvm.sock"),
             name,
             protected,
             temporary_directory,
@@ -278,7 +303,10 @@
             vm_metric: Mutex::new(Default::default()),
             payload_state: Mutex::new(PayloadState::Starting),
             payload_state_updated: Condvar::new(),
-        })
+            requester_uid_name,
+        };
+        info!("{} created", &instance);
+        Ok(instance)
     }
 
     /// Starts an instance of `crosvm` to manage the VM. The `crosvm` instance will be killed when
@@ -286,7 +314,11 @@
     pub fn start(self: &Arc<Self>) -> Result<(), Error> {
         let mut vm_metric = self.vm_metric.lock().unwrap();
         vm_metric.start_timestamp = Some(SystemTime::now());
-        self.vm_state.lock().unwrap().start(self.clone())
+        let ret = self.vm_state.lock().unwrap().start(self.clone());
+        if ret.is_ok() {
+            info!("{} started", &self);
+        }
+        ret.with_context(|| format!("{} failed to start", &self))
     }
 
     /// Monitors the exit of the VM (i.e. termination of the `child` process). When that happens,
@@ -310,6 +342,7 @@
         *vm_state = VmState::Dead;
         // Ensure that the mutex is released before calling the callbacks.
         drop(vm_state);
+        info!("{} exited", &self);
 
         // Read the pipe to see if any failure reason is written
         let mut failure_reason = String::new();
@@ -439,6 +472,46 @@
         }
     }
 
+    /// Responds to memory-trimming notifications by inflating the virtio
+    /// balloon to reclaim guest memory.
+    pub fn trim_memory(&self, level: MemoryTrimLevel) -> Result<(), Error> {
+        let request = VmRequest::BalloonCommand(BalloonControlCommand::Stats {});
+        match vm_control::client::handle_request(&request, &self.crosvm_control_socket_path) {
+            Ok(VmResponse::BalloonStats { stats, balloon_actual: _ }) => {
+                if let Some(total_memory) = stats.total_memory {
+                    // Reclaim up to 50% of total memory assuming worst case
+                    // most memory is anonymous and must be swapped to zram
+                    // with an approximate 2:1 compression ratio.
+                    let pct = match level {
+                        MemoryTrimLevel::TRIM_MEMORY_RUNNING_CRITICAL => 50,
+                        MemoryTrimLevel::TRIM_MEMORY_RUNNING_LOW => 30,
+                        MemoryTrimLevel::TRIM_MEMORY_RUNNING_MODERATE => 10,
+                        _ => bail!("Invalid memory trim level {:?}", level),
+                    };
+                    let command =
+                        BalloonControlCommand::Adjust { num_bytes: total_memory * pct / 100 };
+                    if let Err(e) = vm_control::client::handle_request(
+                        &VmRequest::BalloonCommand(command),
+                        &self.crosvm_control_socket_path,
+                    ) {
+                        bail!("Error sending balloon adjustment: {:?}", e);
+                    }
+                }
+            }
+            Ok(VmResponse::Err(e)) => {
+                // ENOTSUP is returned when the balloon protocol is not initialised. This
+                // can occur for numerous reasons: Guest is still booting, guest doesn't
+                // support ballooning, host doesn't support ballooning. We don't log or
+                // raise an error in this case: trim is just a hint and we can ignore it.
+                if e.errno() != libc::ENOTSUP {
+                    bail!("Errno return when requesting balloon stats: {}", e.errno())
+                }
+            }
+            e => bail!("Error requesting balloon stats: {:?}", e),
+        }
+        Ok(())
+    }
+
     /// Checks if ramdump has been created. If so, send a notification to the user with the handle
     /// to read the ramdump.
     fn handle_ramdump(&self) -> Result<(), Error> {
@@ -578,7 +651,7 @@
 /// Starts an instance of `crosvm` to manage a new VM.
 fn run_vm(
     config: CrosvmConfig,
-    temporary_directory: &Path,
+    crosvm_control_socket_path: &Path,
     failure_pipe_write: File,
 ) -> Result<SharedChild, Error> {
     validate_config(&config)?;
@@ -684,9 +757,8 @@
         command.arg(add_preserved_fd(&mut preserved_fds, kernel));
     }
 
-    let control_server_socket =
-        UnixSeqpacketListener::bind(temporary_directory.join("crosvm.sock"))
-            .context("failed to create control server")?;
+    let control_server_socket = UnixSeqpacketListener::bind(crosvm_control_socket_path)
+        .context("failed to create control server")?;
     command.arg("--socket").arg(add_preserved_fd(&mut preserved_fds, &control_server_socket));
 
     debug!("Preserving FDs {:?}", preserved_fds);
diff --git a/vm/Android.bp b/vm/Android.bp
index 95ef082..b95dca3 100644
--- a/vm/Android.bp
+++ b/vm/Android.bp
@@ -2,8 +2,8 @@
     default_applicable_licenses: ["Android-Apache-2.0"],
 }
 
-rust_binary {
-    name: "vm",
+rust_defaults {
+    name: "vm.defaults",
     crate_name: "vm",
     srcs: ["src/main.rs"],
     edition: "2021",
@@ -25,11 +25,23 @@
         "libvmclient",
         "libzip",
     ],
+}
+
+rust_binary {
+    name: "vm",
+    defaults: ["vm.defaults"],
     apex_available: [
         "com.android.virt",
     ],
 }
 
+rust_test {
+    name: "vm.test",
+    defaults: ["vm.defaults"],
+    test_suites: ["general-tests"],
+    compile_multilib: "first",
+}
+
 sh_binary_host {
     name: "vm_shell",
     src: "vm_shell.sh",
diff --git a/vm/src/main.rs b/vm/src/main.rs
index b5046fb..ca25101 100644
--- a/vm/src/main.rs
+++ b/vm/src/main.rs
@@ -393,3 +393,14 @@
 
     Ok(())
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use clap::IntoApp;
+
+    #[test]
+    fn verify_app() {
+        Opt::into_app().debug_assert();
+    }
+}
diff --git a/vm_payload/src/api.rs b/vm_payload/src/api.rs
index a79c0bb..28b440e 100644
--- a/vm_payload/src/api.rs
+++ b/vm_payload/src/api.rs
@@ -18,7 +18,7 @@
 #![warn(unsafe_op_in_unsafe_fn)]
 
 use android_system_virtualization_payload::aidl::android::system::virtualization::payload::IVmPayloadService::{
-    IVmPayloadService, VM_PAYLOAD_SERVICE_SOCKET_NAME, VM_APK_CONTENTS_PATH};
+    ENCRYPTEDSTORE_MOUNTPOINT, IVmPayloadService, VM_PAYLOAD_SERVICE_SOCKET_NAME, VM_APK_CONTENTS_PATH};
 use anyhow::{ensure, bail, Context, Result};
 use binder::{Strong, unstable_api::{AIBinder, new_spibinder}};
 use lazy_static::lazy_static;
@@ -28,6 +28,7 @@
 use std::ffi::CString;
 use std::fmt::Debug;
 use std::os::raw::{c_char, c_void};
+use std::path::Path;
 use std::ptr;
 use std::sync::{Mutex, atomic::{AtomicBool, Ordering}};
 
@@ -35,6 +36,8 @@
     static ref VM_APK_CONTENTS_PATH_C: CString =
         CString::new(VM_APK_CONTENTS_PATH).expect("CString::new failed");
     static ref PAYLOAD_CONNECTION: Mutex<Option<Strong<dyn IVmPayloadService>>> = Mutex::default();
+    static ref VM_ENCRYPTED_STORAGE_PATH_C: CString =
+        CString::new(ENCRYPTEDSTORE_MOUNTPOINT).expect("CString::new failed");
 }
 
 static ALREADY_NOTIFIED: AtomicBool = AtomicBool::new(false);
@@ -249,12 +252,15 @@
 /// Gets the path to the APK contents.
 #[no_mangle]
 pub extern "C" fn AVmPayload_getApkContentsPath() -> *const c_char {
-    (*VM_APK_CONTENTS_PATH_C).as_ptr()
+    VM_APK_CONTENTS_PATH_C.as_ptr()
 }
 
 /// Gets the path to the VM's encrypted storage.
 #[no_mangle]
 pub extern "C" fn AVmPayload_getEncryptedStoragePath() -> *const c_char {
-    // TODO(b/254454578): Return a real path if storage is present
-    ptr::null()
+    if Path::new(ENCRYPTEDSTORE_MOUNTPOINT).exists() {
+        VM_ENCRYPTED_STORAGE_PATH_C.as_ptr()
+    } else {
+        ptr::null()
+    }
 }