Plumb onTrimMemory app callback to virtualizationservice

Virtualizationservice uses this hint to reclaim guest memory
via the virtio balloon.

Bug: 238931615
Test: atest MicrodroidBenchmarks
Change-Id: Ie4e70160a3305dec68def27667cecd4cd543293f
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..e535258 100644
--- a/virtualizationservice/src/crosvm.rs
+++ b/virtualizationservice/src/crosvm.rs
@@ -38,7 +38,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 +50,7 @@
 
 /// external/crosvm
 use base::UnixSeqpacketListener;
+use vm_control::{BalloonControlCommand, VmRequest, VmResponse};
 
 const CROSVM_PATH: &str = "/apex/com.android.virt/bin/crosvm";
 
@@ -171,7 +175,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 +232,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.
@@ -268,6 +274,7 @@
             vm_state: Mutex::new(VmState::NotStarted { config }),
             vm_context,
             cid,
+            crosvm_control_socket_path: temporary_directory.join("crosvm.sock"),
             name,
             protected,
             temporary_directory,
@@ -439,6 +446,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 +625,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 +731,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);