Send atoms via VirtualizationServiceInternal

In preparation for spliting VS into two services (global and
UID-specific), forward all statsd atoms via
VirtualizationServiceInternal. This is necessary because statsd
will not accept atoms from the client-specific VS instance, as it only
accepts atoms from known UIDs.

Bug: 245727626
Test: atest MicrodroidHostTestCases
Change-Id: I779bbc77a83665a88c9e12d3ad25824e7857514b
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachine.java b/javalib/src/android/system/virtualmachine/VirtualMachine.java
index 1ea6714..db5b4d1 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachine.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachine.java
@@ -61,8 +61,8 @@
 import android.os.RemoteException;
 import android.os.ServiceManager;
 import android.os.ServiceSpecificException;
+import android.system.virtualizationcommon.DeathReason;
 import android.system.virtualizationcommon.ErrorCode;
-import android.system.virtualizationservice.DeathReason;
 import android.system.virtualizationservice.IVirtualMachine;
 import android.system.virtualizationservice.IVirtualMachineCallback;
 import android.system.virtualizationservice.IVirtualizationService;
diff --git a/virtualizationservice/aidl/Android.bp b/virtualizationservice/aidl/Android.bp
index a0bbc00..f028c0f 100644
--- a/virtualizationservice/aidl/Android.bp
+++ b/virtualizationservice/aidl/Android.bp
@@ -36,6 +36,7 @@
 aidl_interface {
     name: "android.system.virtualizationservice_internal",
     srcs: ["android/system/virtualizationservice_internal/**/*.aidl"],
+    imports: ["android.system.virtualizationcommon"],
     unstable: true,
     backend: {
         java: {
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice/DeathReason.aidl b/virtualizationservice/aidl/android/system/virtualizationcommon/DeathReason.aidl
similarity index 97%
rename from virtualizationservice/aidl/android/system/virtualizationservice/DeathReason.aidl
rename to virtualizationservice/aidl/android/system/virtualizationcommon/DeathReason.aidl
index 0980630..3f47002 100644
--- a/virtualizationservice/aidl/android/system/virtualizationservice/DeathReason.aidl
+++ b/virtualizationservice/aidl/android/system/virtualizationcommon/DeathReason.aidl
@@ -13,7 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package android.system.virtualizationservice;
+package android.system.virtualizationcommon;
 
 /**
  * The reason why a VM died.
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualMachineCallback.aidl b/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualMachineCallback.aidl
index 34b6fa5..40e4c58 100644
--- a/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualMachineCallback.aidl
+++ b/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualMachineCallback.aidl
@@ -15,8 +15,8 @@
  */
 package android.system.virtualizationservice;
 
+import android.system.virtualizationcommon.DeathReason;
 import android.system.virtualizationcommon.ErrorCode;
-import android.system.virtualizationservice.DeathReason;
 
 /**
  * An object which a client may register with the VirtualizationService to get callbacks about the
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice_internal/AtomVmBooted.aidl b/virtualizationservice/aidl/android/system/virtualizationservice_internal/AtomVmBooted.aidl
new file mode 100644
index 0000000..2d9e9bf
--- /dev/null
+++ b/virtualizationservice/aidl/android/system/virtualizationservice_internal/AtomVmBooted.aidl
@@ -0,0 +1,23 @@
+/*
+ * 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_internal;
+
+parcelable AtomVmBooted {
+    int uid;
+    @utf8InCpp String vmIdentifier;
+    long elapsedTimeMillis;
+}
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice_internal/AtomVmCreationRequested.aidl b/virtualizationservice/aidl/android/system/virtualizationservice_internal/AtomVmCreationRequested.aidl
new file mode 100644
index 0000000..ec2d206
--- /dev/null
+++ b/virtualizationservice/aidl/android/system/virtualizationservice_internal/AtomVmCreationRequested.aidl
@@ -0,0 +1,29 @@
+/*
+ * 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_internal;
+
+parcelable AtomVmCreationRequested {
+    int uid;
+    @utf8InCpp String vmIdentifier;
+    boolean isProtected;
+    boolean creationSucceeded;
+    int binderExceptionCode;
+    int configType;
+    int numCpus;
+    int memoryMib;
+    @utf8InCpp String apexes;
+}
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice_internal/AtomVmExited.aidl b/virtualizationservice/aidl/android/system/virtualizationservice_internal/AtomVmExited.aidl
new file mode 100644
index 0000000..60b902d
--- /dev/null
+++ b/virtualizationservice/aidl/android/system/virtualizationservice_internal/AtomVmExited.aidl
@@ -0,0 +1,30 @@
+/*
+ * 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_internal;
+
+import android.system.virtualizationcommon.DeathReason;
+
+parcelable AtomVmExited {
+    int uid;
+    String vmIdentifier;
+    DeathReason deathReason;
+    int exitSignal;
+    long elapsedTimeMillis;
+    long guestTimeMillis;
+    long rssVmKb;
+    long rssCrosvmKb;
+}
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice_internal/IVirtualizationServiceInternal.aidl b/virtualizationservice/aidl/android/system/virtualizationservice_internal/IVirtualizationServiceInternal.aidl
index 851ddf4..ff56b68 100644
--- a/virtualizationservice/aidl/android/system/virtualizationservice_internal/IVirtualizationServiceInternal.aidl
+++ b/virtualizationservice/aidl/android/system/virtualizationservice_internal/IVirtualizationServiceInternal.aidl
@@ -15,6 +15,9 @@
  */
 package android.system.virtualizationservice_internal;
 
+import android.system.virtualizationservice_internal.AtomVmBooted;
+import android.system.virtualizationservice_internal.AtomVmCreationRequested;
+import android.system.virtualizationservice_internal.AtomVmExited;
 import android.system.virtualizationservice_internal.IGlobalVmContext;
 
 interface IVirtualizationServiceInternal {
@@ -26,4 +29,13 @@
      * to the returned object.
      */
     IGlobalVmContext allocateGlobalVmContext();
+
+    /** Forwards a VmBooted atom to statsd. */
+    void atomVmBooted(in AtomVmBooted atom);
+
+    /** Forwards a VmCreationRequested atom to statsd. */
+    void atomVmCreationRequested(in AtomVmCreationRequested atom);
+
+    /** Forwards a VmExited atom to statsd. */
+    void atomVmExited(in AtomVmExited atom);
 }
diff --git a/virtualizationservice/src/aidl.rs b/virtualizationservice/src/aidl.rs
index a35c2ac..3e7eca1 100644
--- a/virtualizationservice/src/aidl.rs
+++ b/virtualizationservice/src/aidl.rs
@@ -14,15 +14,19 @@
 
 //! Implementation of the AIDL interface of the VirtualizationService.
 
-use crate::atom::{write_vm_booted_stats, write_vm_creation_stats};
+use crate::atom::{
+    forward_vm_booted_atom, forward_vm_creation_atom, forward_vm_exited_atom,
+    write_vm_booted_stats, write_vm_creation_stats};
 use crate::composite::make_composite_image;
 use crate::crosvm::{CrosvmConfig, DiskFile, PayloadState, VmContext, VmInstance, VmState};
 use crate::payload::{add_microdroid_payload_images, add_microdroid_system_images};
 use crate::selinux::{getfilecon, SeContext};
 use android_os_permissions_aidl::aidl::android::os::IPermissionController;
-use android_system_virtualizationcommon::aidl::android::system::virtualizationcommon::ErrorCode::ErrorCode;
-use android_system_virtualizationservice::aidl::android::system::virtualizationservice::{
+use android_system_virtualizationcommon::aidl::android::system::virtualizationcommon::{
     DeathReason::DeathReason,
+    ErrorCode::ErrorCode,
+};
+use android_system_virtualizationservice::aidl::android::system::virtualizationservice::{
     DiskImage::DiskImage,
     IVirtualMachine::{BnVirtualMachine, IVirtualMachine},
     IVirtualMachineCallback::IVirtualMachineCallback,
@@ -38,6 +42,9 @@
     VirtualMachineState::VirtualMachineState,
 };
 use android_system_virtualizationservice_internal::aidl::android::system::virtualizationservice_internal::{
+    AtomVmBooted::AtomVmBooted,
+    AtomVmCreationRequested::AtomVmCreationRequested,
+    AtomVmExited::AtomVmExited,
     IGlobalVmContext::{BnGlobalVmContext, IGlobalVmContext},
     IVirtualizationServiceInternal::{BnVirtualizationServiceInternal, IVirtualizationServiceInternal},
 };
@@ -51,6 +58,7 @@
     Status, StatusCode, Strong, ThreadState,
 };
 use disk::QcowFile;
+use lazy_static::lazy_static;
 use libc::VMADDR_CID_HOST;
 use log::{debug, error, info, warn};
 use microdroid_payload_config::{OsConfig, Task, TaskType, VmPayloadConfig};
@@ -102,6 +110,13 @@
 
 const UNFORMATTED_STORAGE_MAGIC: &str = "UNFORMATTED-STORAGE";
 
+lazy_static! {
+    pub static ref GLOBAL_SERVICE: Strong<dyn IVirtualizationServiceInternal> = {
+        let service = VirtualizationServiceInternal::init();
+        BnVirtualizationServiceInternal::new_binder(service, BinderFeatures::default())
+    };
+}
+
 fn is_valid_guest_cid(cid: Cid) -> bool {
     (GUEST_CID_MIN..=GUEST_CID_MAX).contains(&cid)
 }
@@ -164,6 +179,21 @@
         })?;
         Ok(GlobalVmContext::create(cid))
     }
+
+    fn atomVmBooted(&self, atom: &AtomVmBooted) -> Result<(), Status> {
+        forward_vm_booted_atom(atom);
+        Ok(())
+    }
+
+    fn atomVmCreationRequested(&self, atom: &AtomVmCreationRequested) -> Result<(), Status> {
+        forward_vm_creation_atom(atom);
+        Ok(())
+    }
+
+    fn atomVmExited(&self, atom: &AtomVmExited) -> Result<(), Status> {
+        forward_vm_exited_atom(atom);
+        Ok(())
+    }
 }
 
 /// The mutable state of the VirtualizationServiceInternal. There should only be one instance
@@ -257,10 +287,9 @@
 }
 
 /// Implementation of `IVirtualizationService`, the entry point of the AIDL service.
-#[derive(Debug)]
+#[derive(Debug, Default)]
 pub struct VirtualizationService {
     state: Arc<Mutex<State>>,
-    global_service: Strong<dyn IVirtualizationServiceInternal>,
 }
 
 impl Interface for VirtualizationService {
@@ -461,18 +490,14 @@
 
 impl VirtualizationService {
     pub fn init() -> VirtualizationService {
-        let global_service = VirtualizationServiceInternal::init();
-        let global_service =
-            BnVirtualizationServiceInternal::new_binder(global_service, BinderFeatures::default());
-
-        VirtualizationService { global_service, state: Default::default() }
+        VirtualizationService::default()
     }
 
     fn create_vm_context(&self) -> Result<(VmContext, Cid)> {
         const NUM_ATTEMPTS: usize = 5;
 
         for _ in 0..NUM_ATTEMPTS {
-            let global_context = self.global_service.allocateGlobalVmContext()?;
+            let global_context = GLOBAL_SERVICE.allocateGlobalVmContext()?;
             let cid = global_context.getCid()? as Cid;
             let service = VirtualMachineService::new_binder(self.state.clone(), cid).as_binder();
 
diff --git a/virtualizationservice/src/atom.rs b/virtualizationservice/src/atom.rs
index a880d60..698dbca 100644
--- a/virtualizationservice/src/atom.rs
+++ b/virtualizationservice/src/atom.rs
@@ -14,15 +14,20 @@
 
 //! Functions for creating and collecting atoms.
 
-use crate::aidl::clone_file;
+use crate::aidl::{clone_file, GLOBAL_SERVICE};
 use crate::crosvm::VmMetric;
+use android_system_virtualizationcommon::aidl::android::system::virtualizationcommon::DeathReason::DeathReason;
 use android_system_virtualizationservice::aidl::android::system::virtualizationservice::{
-    DeathReason::DeathReason,
     IVirtualMachine::IVirtualMachine,
     VirtualMachineAppConfig::{Payload::Payload, VirtualMachineAppConfig},
     VirtualMachineConfig::VirtualMachineConfig,
 };
 use android_system_virtualizationservice::binder::{Status, Strong};
+use android_system_virtualizationservice_internal::aidl::android::system::virtualizationservice_internal::{
+    AtomVmBooted::AtomVmBooted,
+    AtomVmCreationRequested::AtomVmCreationRequested,
+    AtomVmExited::AtomVmExited,
+};
 use anyhow::{anyhow, Result};
 use binder::{ParcelFileDescriptor, ThreadState};
 use log::{trace, warn};
@@ -106,34 +111,58 @@
         ),
     };
 
-    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
-        };
+    let atom = AtomVmCreationRequested {
+        uid: ThreadState::get_calling_uid() as i32,
+        vmIdentifier: vm_identifier,
+        isProtected: is_protected,
+        creationSucceeded: creation_succeeded,
+        binderExceptionCode: binder_exception_code,
+        configType: config_type as i32,
+        numCpus: num_cpus,
+        memoryMib: memory_mib,
+        apexes,
+    };
 
-        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"),
-        }
+    thread::spawn(move || {
+        GLOBAL_SERVICE.atomVmCreationRequested(&atom).unwrap_or_else(|e| {
+            warn!("Failed to write VmCreationRequested atom: {e}");
+        });
     });
 }
 
+pub fn forward_vm_creation_atom(atom: &AtomVmCreationRequested) {
+    let config_type = match atom.configType {
+        x if x == vm_creation_requested::ConfigType::VirtualMachineAppConfig as i32 => {
+            vm_creation_requested::ConfigType::VirtualMachineAppConfig
+        }
+        x if x == vm_creation_requested::ConfigType::VirtualMachineRawConfig as i32 => {
+            vm_creation_requested::ConfigType::VirtualMachineRawConfig
+        }
+        _ => vm_creation_requested::ConfigType::UnknownConfig,
+    };
+    let vm_creation_requested = vm_creation_requested::VmCreationRequested {
+        uid: atom.uid,
+        vm_identifier: &atom.vmIdentifier,
+        hypervisor: vm_creation_requested::Hypervisor::Pkvm,
+        is_protected: atom.isProtected,
+        creation_succeeded: atom.creationSucceeded,
+        binder_exception_code: atom.binderExceptionCode,
+        config_type,
+        num_cpus: atom.numCpus,
+        cpu_affinity: "", // deprecated
+        memory_mib: atom.memoryMib,
+        apexes: &atom.apexes,
+        // TODO(seungjaeyoo) Fill information about task_profile
+        // TODO(seungjaeyoo) Fill information about disk_image for raw config
+    };
+
+    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"),
+    }
+}
+
 /// 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(
@@ -143,22 +172,34 @@
 ) {
     let vm_identifier = vm_identifier.to_owned();
     let duration = get_duration(vm_start_timestamp);
+
+    let atom = AtomVmBooted {
+        uid,
+        vmIdentifier: vm_identifier,
+        elapsedTimeMillis: duration.as_millis() as i64,
+    };
+
     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"),
-        }
+        GLOBAL_SERVICE.atomVmBooted(&atom).unwrap_or_else(|e| {
+            warn!("Failed to write VmCreationRequested atom: {e}");
+        });
     });
 }
 
+pub fn forward_vm_booted_atom(atom: &AtomVmBooted) {
+    let vm_booted = vm_booted::VmBooted {
+        uid: atom.uid,
+        vm_identifier: &atom.vmIdentifier,
+        elapsed_time_millis: atom.elapsedTimeMillis,
+    };
+
+    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"),
+    }
+}
+
 /// 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(
@@ -173,64 +214,82 @@
     let guest_time_millis = vm_metric.cpu_guest_time.unwrap_or_default();
     let rss = vm_metric.rss.unwrap_or_default();
 
+    let atom = AtomVmExited {
+        uid,
+        vmIdentifier: vm_identifier,
+        elapsedTimeMillis: elapsed_time_millis,
+        deathReason: reason,
+        guestTimeMillis: guest_time_millis,
+        rssVmKb: rss.vm,
+        rssCrosvmKb: rss.crosvm,
+        exitSignal: exit_signal.unwrap_or_default(),
+    };
+
     thread::spawn(move || {
-        let vm_exited = vm_exited::VmExited {
-            uid,
-            vm_identifier: &vm_identifier,
-            elapsed_time_millis,
-            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::START_FAILED => 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,
-            },
-            guest_time_millis,
-            rss_vm_kb: rss.vm,
-            rss_crosvm_kb: rss.crosvm,
-            exit_signal: exit_signal.unwrap_or_default(),
-        };
-        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);
-            }
-            Ok(_) => trace!("statslog_rust succeeded for virtualization service"),
-        }
+        GLOBAL_SERVICE.atomVmExited(&atom).unwrap_or_else(|e| {
+            warn!("Failed to write VmExited atom: {e}");
+        });
     });
 }
 
+pub fn forward_vm_exited_atom(atom: &AtomVmExited) {
+    let death_reason = match atom.deathReason {
+        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::START_FAILED => 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,
+    };
+
+    let vm_exited = vm_exited::VmExited {
+        uid: atom.uid,
+        vm_identifier: &atom.vmIdentifier,
+        elapsed_time_millis: atom.elapsedTimeMillis,
+        death_reason,
+        guest_time_millis: atom.guestTimeMillis,
+        rss_vm_kb: atom.rssVmKb,
+        rss_crosvm_kb: atom.rssCrosvmKb,
+        exit_signal: atom.exitSignal,
+    };
+
+    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),
+        Ok(_) => trace!("statslog_rust succeeded for virtualization service"),
+    }
+}
+
 fn wait_for_statsd() -> Result<()> {
     let mut prop = system_properties::PropertyWatcher::new("init.svc.statsd")?;
     loop {
diff --git a/virtualizationservice/src/crosvm.rs b/virtualizationservice/src/crosvm.rs
index 2c6e370..94248f8 100644
--- a/virtualizationservice/src/crosvm.rs
+++ b/virtualizationservice/src/crosvm.rs
@@ -40,10 +40,8 @@
 use std::sync::{Arc, Condvar, Mutex};
 use std::time::{Duration, SystemTime};
 use std::thread::{self, JoinHandle};
-use android_system_virtualizationservice::aidl::android::system::virtualizationservice::{
-    DeathReason::DeathReason,
-    MemoryTrimLevel::MemoryTrimLevel,
-};
+use android_system_virtualizationcommon::aidl::android::system::virtualizationcommon::DeathReason::DeathReason;
+use android_system_virtualizationservice::aidl::android::system::virtualizationservice::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;
diff --git a/vmclient/src/death_reason.rs b/vmclient/src/death_reason.rs
index 34d89fc..c417a7c 100644
--- a/vmclient/src/death_reason.rs
+++ b/vmclient/src/death_reason.rs
@@ -12,9 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use android_system_virtualizationservice::{
-        aidl::android::system::virtualizationservice::{
-            DeathReason::DeathReason as AidlDeathReason}};
+use android_system_virtualizationcommon::aidl::android::system::virtualizationcommon::DeathReason::DeathReason as AidlDeathReason;
 
 /// The reason why a VM died.
 #[derive(Clone, Copy, Debug, Eq, PartialEq)]
diff --git a/vmclient/src/lib.rs b/vmclient/src/lib.rs
index 7c05545..2f22316 100644
--- a/vmclient/src/lib.rs
+++ b/vmclient/src/lib.rs
@@ -23,10 +23,11 @@
 pub use crate::error_code::ErrorCode;
 pub use crate::errors::VmWaitError;
 use crate::sync::Monitor;
-use android_system_virtualizationcommon::aidl::android::system::virtualizationcommon::ErrorCode::ErrorCode as AidlErrorCode;
+use android_system_virtualizationcommon::aidl::android::system::virtualizationcommon::{
+    DeathReason::DeathReason as AidlDeathReason, ErrorCode::ErrorCode as AidlErrorCode,
+};
 use android_system_virtualizationservice::{
     aidl::android::system::virtualizationservice::{
-        DeathReason::DeathReason as AidlDeathReason,
         IVirtualMachine::IVirtualMachine,
         IVirtualMachineCallback::{BnVirtualMachineCallback, IVirtualMachineCallback},
         IVirtualizationService::IVirtualizationService,