Implement write_vm_exited_stats for VmExited atom

Bug: 236252957
Test: N/A

Change-Id: I167148cd49b6985aa7a19a800615c83f823e0eec
diff --git a/virtualizationservice/src/atom.rs b/virtualizationservice/src/atom.rs
index 82ff161..7f1d6b1 100644
--- a/virtualizationservice/src/atom.rs
+++ b/virtualizationservice/src/atom.rs
@@ -16,14 +16,14 @@
 
 use crate::aidl::clone_file;
 use android_system_virtualizationservice::aidl::android::system::virtualizationservice::{
-    IVirtualMachine::IVirtualMachine, VirtualMachineAppConfig::VirtualMachineAppConfig,
-    VirtualMachineConfig::VirtualMachineConfig,
+    DeathReason::DeathReason, IVirtualMachine::IVirtualMachine,
+    VirtualMachineAppConfig::VirtualMachineAppConfig, VirtualMachineConfig::VirtualMachineConfig,
 };
 use android_system_virtualizationservice::binder::{Status, Strong};
 use anyhow::{anyhow, Result};
 use log::{trace, warn};
 use microdroid_payload_config::VmPayloadConfig;
-use statslog_virtualization_rust::{vm_booted, vm_creation_requested};
+use statslog_virtualization_rust::{vm_booted, vm_creation_requested, vm_exited};
 use zip::ZipArchive;
 
 fn get_vm_payload_config(config: &VirtualMachineAppConfig) -> Result<VmPayloadConfig> {
@@ -128,3 +128,57 @@
         Ok(_) => trace!("statslog_rust succeeded for virtualization service"),
     }
 }
+
+/// Write the stats of VM exit to statsd
+pub fn write_vm_exited_stats(reason: DeathReason) {
+    let empty_string = String::new();
+    let vm_exited = vm_exited::VmExited {
+        // TODO(seungjaeyoo) Implement sending proper data about uid & vm_identifier
+        uid: -1,
+        vm_identifier: &empty_string,
+        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,
+        },
+    };
+    match vm_exited.stats_write() {
+        Err(e) => {
+            warn!("statslog_rust failed with error: {}", e);
+        }
+        Ok(_) => trace!("statslog_rust succeeded for virtualization service"),
+    }
+}
diff --git a/virtualizationservice/src/crosvm.rs b/virtualizationservice/src/crosvm.rs
index 3e1a151..b4d173c 100644
--- a/virtualizationservice/src/crosvm.rs
+++ b/virtualizationservice/src/crosvm.rs
@@ -15,6 +15,7 @@
 //! Functions for running instances of `crosvm`.
 
 use crate::aidl::VirtualMachineCallbacks;
+use crate::atom::write_vm_exited_stats;
 use crate::Cid;
 use anyhow::{bail, Context, Error};
 use command_fds::CommandFdExt;
@@ -260,7 +261,10 @@
             };
 
         self.handle_ramdump().unwrap_or_else(|e| error!("Error handling ramdump: {}", e));
-        self.callbacks.callback_on_died(self.cid, death_reason(&result, &failure_reason));
+
+        let death_reason = death_reason(&result, &failure_reason);
+        self.callbacks.callback_on_died(self.cid, death_reason);
+        write_vm_exited_stats(death_reason);
 
         // Delete temporary files.
         if let Err(e) = remove_dir_all(&self.temporary_directory) {