shadercached/src/service/mod.rs - third_party/platform2 - Git at Google

 // Copyright 2023 The ChromiumOS Authors
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.

 // service module provides code entry path from D-BUS. Hence, functions here
 // are naturally mapped to D-BUS methods.
 // D-BUS methods and signals for other services (ex. concierge) are in
 // individual modules.

 mod concierge;
 mod dlc;
 pub mod helper;
 pub mod signal;
 pub mod spaced;

 use crate::common::*;
 use crate::dbus_wrapper::DbusConnectionTrait;
 use crate::dlc_queue::DlcQueuePtr;
 use crate::shader_cache_mount::{ShaderCacheMount, ShaderCacheMountMapPtr, VmId};

 use anyhow::{anyhow, Result};
 use dbus::MethodErr;
 use log::{debug, error, info, warn};
 use protobuf::Message;
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
 use system_api::shadercached::{
     InstallRequest, InstallResponse, PrepareShaderCacheRequest, PrepareShaderCacheResponse,
     PurgeRequest, UninstallRequest, UnmountRequest,
 };

 // Selectively expose service methods
 pub use concierge::add_shader_cache_group_permission;
 pub use concierge::handle_vm_stopped;
 pub use dlc::handle_dlc_state_changed;
 pub use dlc::mount_dlc;
 pub use dlc::periodic_dlc_handler;
 pub use spaced::handle_disk_space_update;

 #[cfg(test)]
 use helper::unsafe_ops::mock_quota::set_quota_normal;
 #[cfg(not(test))]
 use helper::unsafe_ops::quota::set_quota_normal;

 pub async fn handle_install<D: DbusConnectionTrait>(
     raw_bytes: Vec<u8>,
     mount_map: ShaderCacheMountMapPtr,
     dlc_queue: DlcQueuePtr,
     dbus_conn: Arc<D>,
 ) -> Result<std::vec::Vec<u8>> {
     let request: InstallRequest = Message::parse_from_bytes(&raw_bytes)?;

     let vm_id = VmId {
         vm_name: request.vm_name,
         vm_owner_id: request.vm_owner_id,
     };

     let mut mut_mount_map = mount_map.write().await;
     // TODO(b/271776528): move the insert ShaderCacheMount to
     // PrepareShaderCache method response once we support local RW shader
     // cache in shadercached.
     if !mut_mount_map.contains_key(&vm_id) {
         let vm_gpu_cache_path =
             Path::new(&concierge::get_vm_gpu_cache_path(&vm_id, dbus_conn.clone()).await?)
                 .to_path_buf();
         mut_mount_map.insert(
             vm_id.clone(),
             ShaderCacheMount::new(vm_gpu_cache_path, &vm_id)?,
         );
     } else {
         debug!("Reusing mount information for {:?}", vm_id);
     }

     let shader_cache_mount = mut_mount_map
         .get_mut(&vm_id)
         .ok_or_else(|| MethodErr::failed("Failed to get mount information"))?;

     // Mesa cache path initialization must succeed before we enqueue mount
     // Repeated initializations are no-op.
     // |initialize()| will fail iff mesa shader cache has not been
     // initialized - i.e. VM launched but no applications have run before. Under
     // Normal circumstances, that's not happening because Steam Store UI would
     // have triggered mesa to initialize shader cache.
     // Hence, treat initialize attempt failure as failure. We are not supporting
     // the case of manually triggering shadercached before Steam launch.
     shader_cache_mount
         .initialize(&vm_id, dbus_conn.clone())
         .await?;

     // Save the current mount status in the response
     let mut response: InstallResponse = InstallResponse::new();
     if let Ok(is_mounted) = shader_cache_mount.is_game_mounted(request.steam_app_id, None) {
         // Shader cache might be mounted from the previous mount call.
         response.mounted = is_mounted;
     }

     if request.mount && !response.mounted {
         // Queue mount if not mounted already
         shader_cache_mount.enqueue_mount(request.steam_app_id);
     }

     if !is_dlc_installed(request.steam_app_id) {
         if response.mounted {
             error!("Invalid state, DLC not installed but reported as mount");
             error!("Proceeding with DLC installation");
         }
         // Queue install and upon installation completion, mount if queued.
         info!("Acquiring dlc write mutex to queue install");
         let mut dlc_queue = dlc_queue.write().await;
         if let Some(dequeued_game) = dlc_queue.queue_install(&request.steam_app_id) {
             shader_cache_mount.dequeue_mount(&dequeued_game);
         }
         info!("Successfully queued installation");
         return Ok(response.write_to_bytes()?);
     }

     drop(mut_mount_map);
     debug!(
         "Shader cached already installed for {}",
         request.steam_app_id
     );
     if request.mount && !response.mounted {
         // Mount if dlc is not mounted already
         response.mounted = mount_dlc(request.steam_app_id, mount_map.clone(), dbus_conn.clone())
             .await
             .is_ok();
         debug!("Mounting result: {}", response.mounted);
     }
     Ok(response.write_to_bytes()?)
 }

 fn is_dlc_installed(steam_app_id: SteamAppId) -> bool {
     // /run/imageloader/<dlc_name>/package must exist for DLC to be
     // installed.
     IMAGE_LOADER
         .join(steam_app_id_to_dlc(steam_app_id))
         .join("package")
         .exists()
 }

 pub async fn handle_uninstall(raw_bytes: Vec<u8>, dlc_queue: DlcQueuePtr) -> Result<()> {
     let request: UninstallRequest = protobuf::Message::parse_from_bytes(&raw_bytes)?;
     let mut dlc_queue = dlc_queue.write().await;
     dlc_queue.queue_uninstall(&request.steam_app_id);
     Ok(())
 }

 pub async fn unmount_and_uninstall_all_shader_cache_dlcs<D: DbusConnectionTrait>(
     mount_map: ShaderCacheMountMapPtr,
     dbus_conn: Arc<D>,
 ) -> Result<()> {
     mount_map.clear_all_mounts(None).await?;
     mount_map
         .wait_unmount_completed(None, UNMOUNTER_INTERVAL * 2)
         .await?;

     // TODO(b/270262568): Queue DLC uninstallations instead of waiting for
     // unmounts. Ensure to disallow mounts during the period and allow mounts
     // after completed.
     dlc::uninstall_all_shader_cache_dlcs(dbus_conn.clone()).await?;

     Ok(())
 }

 pub async fn handle_purge<D: DbusConnectionTrait>(
     raw_bytes: Vec<u8>,
     mount_map: ShaderCacheMountMapPtr,
     dbus_conn: Arc<D>,
 ) -> Result<()> {
     let request: PurgeRequest = protobuf::Message::parse_from_bytes(&raw_bytes)?;

     let mount_map_size = { mount_map.clone().read().await.len() };
     if mount_map_size <= 1 {
         // If only one active user, unmount for all users and uninstall all
         // DLCs. Shadercached does not keep state of users globally, hence
         // it cannot know if Borealis is still in use for other users. This
         // means another user may have to redownload the DLC, which is fine for
         // now.
         unmount_and_uninstall_all_shader_cache_dlcs(mount_map.clone(), dbus_conn.clone()).await?;
     }
     let vm_id = VmId {
         vm_name: request.vm_name,
         vm_owner_id: request.vm_owner_id,
     };
     let cleared_dirs = spaced::delete_precompiled_cache(mount_map.clone(), vm_id).await?;
     for dir_path in cleared_dirs {
         std::fs::remove_dir(dir_path)?;
     }
     Ok(())
 }

 pub async fn handle_prepare_shader_cache(
     raw_bytes: Vec<u8>,
     _mount_map: ShaderCacheMountMapPtr,
 ) -> Result<std::vec::Vec<u8>> {
     let request: PrepareShaderCacheRequest = protobuf::Message::parse_from_bytes(&raw_bytes)?;
     let mut response = PrepareShaderCacheResponse::new();
     if request.vm_name.is_empty() || request.vm_owner_id.is_empty() {
         return Err(anyhow!("vm_name and vm_owner_id must be set"));
     }

     let vm_id = VmId {
         vm_name: request.vm_name,
         vm_owner_id: request.vm_owner_id,
     };
     debug!("Preparing shader cache for {:?}", vm_id);

     // TODO(b/271776528): insert ShaderCacheMount here once we support local RW
     // shader cache in shadercached.
     let shader_cache_mount = ShaderCacheMount::new(PathBuf::new(), &vm_id)?;
     let path_str = shader_cache_mount.local_precompiled_cache_path()?;
     let path = Path::new(&path_str);
     if !path.exists() {
         std::fs::create_dir_all(&path_str)?;
     }
     if let Err(e) = set_quota_normal(path) {
         warn!("Failed to set quota for {}: {}", path_str, e);
     }
     response.precompiled_cache_path = path_str;

     Ok(response.write_to_bytes()?)
 }

 pub async fn handle_unmount(raw_bytes: Vec<u8>, mount_map: ShaderCacheMountMapPtr) -> Result<()> {
     let request: UnmountRequest = protobuf::Message::parse_from_bytes(&raw_bytes)?;
     let vm_id = VmId {
         vm_name: request.vm_name,
         vm_owner_id: request.vm_owner_id,
     };

     let mut mount_map = mount_map.write().await;
     if let Some(shader_cache_mount) = mount_map.get_mut(&vm_id) {
         // If VM had mounted something before but nothing is currently mounted,
         // unmount will be queued and unmount will succeed because there is
         // nothing mounted.
         // Shadercached will fire unmount signal in result, which can be
         // picked up by `garcon <..> --unmount --wait` process and exit
         // with success code.
         // (note: --wait flag probably only needed for tasts and manual
         // debugging).
         shader_cache_mount.remove_game_from_db_list(request.steam_app_id)?;
     } else {
         return Err(anyhow!("VM had never mounted shader cache"));
     }

     Ok(())
 }
	// Copyright 2023 The ChromiumOS Authors
	// Use of this source code is governed by a BSD-style license that can be
	// found in the LICENSE file.

	// service module provides code entry path from D-BUS. Hence, functions here
	// are naturally mapped to D-BUS methods.
	// D-BUS methods and signals for other services (ex. concierge) are in
	// individual modules.

	mod concierge;
	mod dlc;
	pub mod helper;
	pub mod signal;
	pub mod spaced;

	use crate::common::*;
	use crate::dbus_wrapper::DbusConnectionTrait;
	use crate::dlc_queue::DlcQueuePtr;
	use crate::shader_cache_mount::{ShaderCacheMount, ShaderCacheMountMapPtr, VmId};

	use anyhow::{anyhow, Result};
	use dbus::MethodErr;
	use log::{debug, error, info, warn};
	use protobuf::Message;
	use std::path::{Path, PathBuf};
	use std::sync::Arc;
	use system_api::shadercached::{
	InstallRequest, InstallResponse, PrepareShaderCacheRequest, PrepareShaderCacheResponse,
	PurgeRequest, UninstallRequest, UnmountRequest,
	};

	// Selectively expose service methods
	pub use concierge::add_shader_cache_group_permission;
	pub use concierge::handle_vm_stopped;
	pub use dlc::handle_dlc_state_changed;
	pub use dlc::mount_dlc;
	pub use dlc::periodic_dlc_handler;
	pub use spaced::handle_disk_space_update;

	#[cfg(test)]
	use helper::unsafe_ops::mock_quota::set_quota_normal;
	#[cfg(not(test))]
	use helper::unsafe_ops::quota::set_quota_normal;

	pub async fn handle_install<D: DbusConnectionTrait>(
	raw_bytes: Vec<u8>,
	mount_map: ShaderCacheMountMapPtr,
	dlc_queue: DlcQueuePtr,
	dbus_conn: Arc<D>,
	) -> Result<std::vec::Vec<u8>> {
	let request: InstallRequest = Message::parse_from_bytes(&raw_bytes)?;

	let vm_id = VmId {
	vm_name: request.vm_name,
	vm_owner_id: request.vm_owner_id,
	};

	let mut mut_mount_map = mount_map.write().await;
	// TODO(b/271776528): move the insert ShaderCacheMount to
	// PrepareShaderCache method response once we support local RW shader
	// cache in shadercached.
	if !mut_mount_map.contains_key(&vm_id) {
	let vm_gpu_cache_path =
	Path::new(&concierge::get_vm_gpu_cache_path(&vm_id, dbus_conn.clone()).await?)
	.to_path_buf();
	mut_mount_map.insert(
	vm_id.clone(),
	ShaderCacheMount::new(vm_gpu_cache_path, &vm_id)?,
	);
	} else {
	debug!("Reusing mount information for {:?}", vm_id);
	}

	let shader_cache_mount = mut_mount_map
	.get_mut(&vm_id)
	.ok_or_else(\|\| MethodErr::failed("Failed to get mount information"))?;

	// Mesa cache path initialization must succeed before we enqueue mount
	// Repeated initializations are no-op.
	// \|initialize()\| will fail iff mesa shader cache has not been
	// initialized - i.e. VM launched but no applications have run before. Under
	// Normal circumstances, that's not happening because Steam Store UI would
	// have triggered mesa to initialize shader cache.
	// Hence, treat initialize attempt failure as failure. We are not supporting
	// the case of manually triggering shadercached before Steam launch.
	shader_cache_mount
	.initialize(&vm_id, dbus_conn.clone())
	.await?;

	// Save the current mount status in the response
	let mut response: InstallResponse = InstallResponse::new();
	if let Ok(is_mounted) = shader_cache_mount.is_game_mounted(request.steam_app_id, None) {
	// Shader cache might be mounted from the previous mount call.
	response.mounted = is_mounted;
	}

	if request.mount && !response.mounted {
	// Queue mount if not mounted already
	shader_cache_mount.enqueue_mount(request.steam_app_id);
	}

	if !is_dlc_installed(request.steam_app_id) {
	if response.mounted {
	error!("Invalid state, DLC not installed but reported as mount");
	error!("Proceeding with DLC installation");
	}
	// Queue install and upon installation completion, mount if queued.
	info!("Acquiring dlc write mutex to queue install");
	let mut dlc_queue = dlc_queue.write().await;
	if let Some(dequeued_game) = dlc_queue.queue_install(&request.steam_app_id) {
	shader_cache_mount.dequeue_mount(&dequeued_game);
	}
	info!("Successfully queued installation");
	return Ok(response.write_to_bytes()?);
	}

	drop(mut_mount_map);
	debug!(
	"Shader cached already installed for {}",
	request.steam_app_id
	);
	if request.mount && !response.mounted {
	// Mount if dlc is not mounted already
	response.mounted = mount_dlc(request.steam_app_id, mount_map.clone(), dbus_conn.clone())
	.await
	.is_ok();
	debug!("Mounting result: {}", response.mounted);
	}
	Ok(response.write_to_bytes()?)
	}

	fn is_dlc_installed(steam_app_id: SteamAppId) -> bool {
	// /run/imageloader/<dlc_name>/package must exist for DLC to be
	// installed.
	IMAGE_LOADER
	.join(steam_app_id_to_dlc(steam_app_id))
	.join("package")
	.exists()
	}

	pub async fn handle_uninstall(raw_bytes: Vec<u8>, dlc_queue: DlcQueuePtr) -> Result<()> {
	let request: UninstallRequest = protobuf::Message::parse_from_bytes(&raw_bytes)?;
	let mut dlc_queue = dlc_queue.write().await;
	dlc_queue.queue_uninstall(&request.steam_app_id);
	Ok(())
	}

	pub async fn unmount_and_uninstall_all_shader_cache_dlcs<D: DbusConnectionTrait>(
	mount_map: ShaderCacheMountMapPtr,
	dbus_conn: Arc<D>,
	) -> Result<()> {
	mount_map.clear_all_mounts(None).await?;
	mount_map
	.wait_unmount_completed(None, UNMOUNTER_INTERVAL * 2)
	.await?;

	// TODO(b/270262568): Queue DLC uninstallations instead of waiting for
	// unmounts. Ensure to disallow mounts during the period and allow mounts
	// after completed.
	dlc::uninstall_all_shader_cache_dlcs(dbus_conn.clone()).await?;

	Ok(())
	}

	pub async fn handle_purge<D: DbusConnectionTrait>(
	raw_bytes: Vec<u8>,
	mount_map: ShaderCacheMountMapPtr,
	dbus_conn: Arc<D>,
	) -> Result<()> {
	let request: PurgeRequest = protobuf::Message::parse_from_bytes(&raw_bytes)?;

	let mount_map_size = { mount_map.clone().read().await.len() };
	if mount_map_size <= 1 {
	// If only one active user, unmount for all users and uninstall all
	// DLCs. Shadercached does not keep state of users globally, hence
	// it cannot know if Borealis is still in use for other users. This
	// means another user may have to redownload the DLC, which is fine for
	// now.
	unmount_and_uninstall_all_shader_cache_dlcs(mount_map.clone(), dbus_conn.clone()).await?;
	}
	let vm_id = VmId {
	vm_name: request.vm_name,
	vm_owner_id: request.vm_owner_id,
	};
	let cleared_dirs = spaced::delete_precompiled_cache(mount_map.clone(), vm_id).await?;
	for dir_path in cleared_dirs {
	std::fs::remove_dir(dir_path)?;
	}
	Ok(())
	}

	pub async fn handle_prepare_shader_cache(
	raw_bytes: Vec<u8>,
	_mount_map: ShaderCacheMountMapPtr,
	) -> Result<std::vec::Vec<u8>> {
	let request: PrepareShaderCacheRequest = protobuf::Message::parse_from_bytes(&raw_bytes)?;
	let mut response = PrepareShaderCacheResponse::new();
	if request.vm_name.is_empty() \|\| request.vm_owner_id.is_empty() {
	return Err(anyhow!("vm_name and vm_owner_id must be set"));
	}

	let vm_id = VmId {
	vm_name: request.vm_name,
	vm_owner_id: request.vm_owner_id,
	};
	debug!("Preparing shader cache for {:?}", vm_id);

	// TODO(b/271776528): insert ShaderCacheMount here once we support local RW
	// shader cache in shadercached.
	let shader_cache_mount = ShaderCacheMount::new(PathBuf::new(), &vm_id)?;
	let path_str = shader_cache_mount.local_precompiled_cache_path()?;
	let path = Path::new(&path_str);
	if !path.exists() {
	std::fs::create_dir_all(&path_str)?;
	}
	if let Err(e) = set_quota_normal(path) {
	warn!("Failed to set quota for {}: {}", path_str, e);
	}
	response.precompiled_cache_path = path_str;

	Ok(response.write_to_bytes()?)
	}

	pub async fn handle_unmount(raw_bytes: Vec<u8>, mount_map: ShaderCacheMountMapPtr) -> Result<()> {
	let request: UnmountRequest = protobuf::Message::parse_from_bytes(&raw_bytes)?;
	let vm_id = VmId {
	vm_name: request.vm_name,
	vm_owner_id: request.vm_owner_id,
	};

	let mut mount_map = mount_map.write().await;
	if let Some(shader_cache_mount) = mount_map.get_mut(&vm_id) {
	// If VM had mounted something before but nothing is currently mounted,
	// unmount will be queued and unmount will succeed because there is
	// nothing mounted.
	// Shadercached will fire unmount signal in result, which can be
	// picked up by `garcon <..> --unmount --wait` process and exit
	// with success code.
	// (note: --wait flag probably only needed for tasts and manual
	// debugging).
	shader_cache_mount.remove_game_from_db_list(request.steam_app_id)?;
	} else {
	return Err(anyhow!("VM had never mounted shader cache"));
	}

	Ok(())
	}