Improve write_vm_creation_stats for expanded VmCreationRequested

Bug: 236253343
Test: N/A

Change-Id: Ie81361bef8937296062b4fe2415d20adf44c0ef0
diff --git a/virtualizationservice/src/aidl.rs b/virtualizationservice/src/aidl.rs
index bba75ac..5bc646f 100644
--- a/virtualizationservice/src/aidl.rs
+++ b/virtualizationservice/src/aidl.rs
@@ -14,6 +14,7 @@
 
 //! Implementation of the AIDL interface of the VirtualizationService.
 
+use crate::atom::write_vm_creation_stats;
 use crate::composite::make_composite_image;
 use crate::crosvm::{CrosvmConfig, DiskFile, PayloadState, VmInstance, VmState};
 use crate::payload::add_microdroid_images;
@@ -48,11 +49,10 @@
 use binder_common::rpc_server::run_rpc_server_with_factory;
 use disk::QcowFile;
 use idsig::{HashAlgorithm, V4Signature};
-use log::{debug, error, info, warn, trace};
+use log::{debug, error, info, warn};
 use microdroid_payload_config::VmPayloadConfig;
 use rustutils::system_properties;
 use semver::VersionReq;
-use statslog_virtualization_rust::vm_creation_requested::{stats_write, Hypervisor};
 use std::convert::TryInto;
 use std::ffi::CStr;
 use std::fs::{create_dir, File, OpenOptions};
@@ -131,23 +131,7 @@
     ) -> binder::Result<Strong<dyn IVirtualMachine>> {
         let mut is_protected = false;
         let ret = self.create_vm_internal(config, console_fd, log_fd, &mut is_protected);
-        match ret {
-            Ok(_) => {
-                let ok_status = Status::ok();
-                write_vm_creation_stats(
-                    is_protected,
-                    /*creation_succeeded*/ true,
-                    ok_status.exception_code() as i32,
-                );
-            }
-            Err(ref e) => {
-                write_vm_creation_stats(
-                    is_protected,
-                    /*creation_succeeded*/ false,
-                    e.exception_code() as i32,
-                );
-            }
-        }
+        write_vm_creation_stats(config, is_protected, &ret);
         ret
     }
 
@@ -506,16 +490,6 @@
     }
 }
 
-/// Write the stats of VMCreation to statsd
-fn write_vm_creation_stats(is_protected: bool, creation_succeeded: bool, exception_code: i32) {
-    match stats_write(Hypervisor::Pkvm, is_protected, creation_succeeded, exception_code) {
-        Err(e) => {
-            warn!("statslog_rust failed with error: {}", e);
-        }
-        Ok(_) => trace!("statslog_rust succeeded for virtualization service"),
-    }
-}
-
 /// Waits for incoming connections from VM. If a new connection is made, stores the stream in the
 /// corresponding `VmInstance`.
 fn handle_stream_connection_from_vm(state: Arc<Mutex<State>>) -> Result<()> {
@@ -1003,7 +977,7 @@
 }
 
 /// Converts a `&ParcelFileDescriptor` to a `File` by cloning the file.
-fn clone_file(file: &ParcelFileDescriptor) -> Result<File, Status> {
+pub fn clone_file(file: &ParcelFileDescriptor) -> Result<File, Status> {
     file.as_ref().try_clone().map_err(|e| {
         Status::new_exception_str(
             ExceptionCode::BAD_PARCELABLE,
diff --git a/virtualizationservice/src/atom.rs b/virtualizationservice/src/atom.rs
new file mode 100644
index 0000000..960eaa7
--- /dev/null
+++ b/virtualizationservice/src/atom.rs
@@ -0,0 +1,114 @@
+// 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.
+
+//! Functions for creating and collecting atoms.
+
+use crate::aidl::clone_file;
+use android_system_virtualizationservice::aidl::android::system::virtualizationservice::{
+    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_creation_requested;
+use zip::ZipArchive;
+
+fn get_vm_payload_config(config: &VirtualMachineAppConfig) -> Result<VmPayloadConfig> {
+    let apk = config.apk.as_ref().ok_or_else(|| anyhow!("APK is none"))?;
+    let apk_file = clone_file(apk)?;
+    let mut apk_zip = ZipArchive::new(&apk_file)?;
+    let config_file = apk_zip.by_name(&config.configPath)?;
+    let vm_payload_config: VmPayloadConfig = serde_json::from_reader(config_file)?;
+    Ok(vm_payload_config)
+}
+
+/// Write the stats of VMCreation to statsd
+pub fn write_vm_creation_stats(
+    config: &VirtualMachineConfig,
+    is_protected: bool,
+    ret: &binder::Result<Strong<dyn IVirtualMachine>>,
+) {
+    let creation_succeeded;
+    let binder_exception_code;
+    match ret {
+        Ok(_) => {
+            creation_succeeded = true;
+            binder_exception_code = Status::ok().exception_code() as i32;
+        }
+        Err(ref e) => {
+            creation_succeeded = false;
+            binder_exception_code = e.exception_code() as i32;
+        }
+    }
+
+    let config_type;
+    let num_cpus;
+    let cpu_affinity;
+    let memory_mib;
+    let apexes;
+    match config {
+        VirtualMachineConfig::AppConfig(config) => {
+            config_type = vm_creation_requested::ConfigType::VirtualMachineAppConfig;
+            num_cpus = config.numCpus;
+            cpu_affinity = config.cpuAffinity.clone().unwrap_or_default();
+            memory_mib = config.memoryMib;
+
+            let vm_payload_config = get_vm_payload_config(config);
+            if let Ok(vm_payload_config) = vm_payload_config {
+                apexes = vm_payload_config
+                    .apexes
+                    .iter()
+                    .map(|x| x.name.clone())
+                    .collect::<Vec<String>>()
+                    .join(":");
+            } else {
+                apexes = "INFO: Can't get VmPayloadConfig".into();
+            }
+        }
+        VirtualMachineConfig::RawConfig(config) => {
+            config_type = vm_creation_requested::ConfigType::VirtualMachineRawConfig;
+            num_cpus = config.numCpus;
+            cpu_affinity = config.cpuAffinity.clone().unwrap_or_default();
+            memory_mib = config.memoryMib;
+            apexes = String::new();
+        }
+    }
+
+    let empty_string = String::new();
+    let vm_creation_requested = vm_creation_requested::VmCreationRequested {
+        // TODO(seungjaeyoo) Implement sending proper data about uid & vm_identifier
+        uid: -1,
+        vm_identifier: &empty_string,
+        hypervisor: vm_creation_requested::Hypervisor::Pkvm,
+        is_protected,
+        creation_succeeded,
+        binder_exception_code,
+        config_type,
+        num_cpus,
+        cpu_affinity: &cpu_affinity,
+        memory_mib,
+        apexes: &apexes,
+        // TODO(seungjaeyoo) Fill information about task_profile
+        // TODO(seungjaeyoo) Fill information about disk_image for raw config
+    };
+
+    match vm_creation_requested.stats_write() {
+        Err(e) => {
+            warn!("statslog_rust failed with error: {}", e);
+        }
+        Ok(_) => trace!("statslog_rust succeeded for virtualization service"),
+    }
+}
diff --git a/virtualizationservice/src/main.rs b/virtualizationservice/src/main.rs
index 3b0adb9..93a5966 100644
--- a/virtualizationservice/src/main.rs
+++ b/virtualizationservice/src/main.rs
@@ -15,6 +15,7 @@
 //! Android VirtualizationService
 
 mod aidl;
+mod atom;
 mod composite;
 mod crosvm;
 mod payload;