feat: session restore!

2024-02-20 03:02:57 -05:00
parent c41179a437
commit 8e317a8e08
6 changed files with 242 additions and 81 deletions
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1648,7 +1648,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "7f4c021e1093a56626774e81216a4ce732a735e5bad4868a03f3ed65ca0c3919"
 dependencies = [
 "once_cell",
- "toml_edit",
+ "toml_edit 0.19.15",
 ]

 [[package]]
@@ -1974,6 +1974,15 @@ dependencies = [
 "syn 2.0.37",
 ]

+[[package]]
+name = "serde_spanned"
+version = "0.6.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "eb3622f419d1296904700073ea6cc23ad690adbd66f13ea683df73298736f0c1"
+dependencies = [
+ "serde",
+]
+
 [[package]]
 name = "sharded-slab"
 version = "0.1.4"
@@ -2149,6 +2158,7 @@ dependencies = [
 "stardust-xr-server-codegen",
 "stereokit",
 "tokio",
+ "toml",
 "tracing",
 "tracing-subscriber",
 "tracing-tracy",
@@ -2347,10 +2357,25 @@ dependencies = [
 ]

 [[package]]
-name = "toml_datetime"
-version = "0.6.3"
+name = "toml"
+version = "0.8.10"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7cda73e2f1397b1262d6dfdcef8aafae14d1de7748d66822d3bfeeb6d03e5e4b"
+checksum = "9a9aad4a3066010876e8dcf5a8a06e70a558751117a145c6ce2b82c2e2054290"
+dependencies = [
+ "serde",
+ "serde_spanned",
+ "toml_datetime",
+ "toml_edit 0.22.6",
+]
+
+[[package]]
+name = "toml_datetime"
+version = "0.6.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3550f4e9685620ac18a50ed434eb3aec30db8ba93b0287467bca5826ea25baf1"
+dependencies = [
+ "serde",
+]

 [[package]]
 name = "toml_edit"
@@ -2360,7 +2385,20 @@ checksum = "1b5bb770da30e5cbfde35a2d7b9b8a2c4b8ef89548a7a6aeab5c9a576e3e7421"
 dependencies = [
 "indexmap 2.0.0",
 "toml_datetime",
- "winnow",
+ "winnow 0.5.15",
+]
+
+[[package]]
+name = "toml_edit"
+version = "0.22.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2c1b5fd4128cc8d3e0cb74d4ed9a9cc7c7284becd4df68f5f940e1ad123606f6"
+dependencies = [
+ "indexmap 2.0.0",
+ "serde",
+ "serde_spanned",
+ "toml_datetime",
+ "winnow 0.6.2",
 ]

 [[package]]
@@ -2779,6 +2817,15 @@ dependencies = [
 "memchr",
 ]

+[[package]]
+name = "winnow"
+version = "0.6.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7a4191c47f15cc3ec71fcb4913cb83d58def65dd3787610213c649283b5ce178"
+dependencies = [
+ "memchr",
+]
+
 [[package]]
 name = "x11rb"
 version = "0.13.0"
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -72,6 +72,7 @@ wayland-scanner = "0.31.1"
 wayland-backend = "0.3.3"
 cluFlock = "1.2.7"
 fxtypemap = "0.2.0"
+toml = "0.8.10"

 [dependencies.smithay]
 # git = "https://github.com/technobaboo/smithay.git" # Until we get stereokit to understand OES samplers and external textures
--- a/src/core/client.rs
+++ b/src/core/client.rs
@@ -16,7 +16,7 @@ use stardust_xr::{
 	messenger::{self, MessageSenderHandle},
 	schemas::flex::serialize,
 };
-use std::{fs, iter::FromIterator, path::PathBuf, sync::Arc};
+use std::{fmt::Debug, fs, iter::FromIterator, path::PathBuf, sync::Arc};
 use tokio::{net::UnixStream, task::JoinHandle};
 use tracing::info;

@@ -215,6 +215,19 @@ impl Client {
 		}
 	}
 }
+impl Debug for Client {
+	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+		f.debug_struct("Client")
+			.field("pid", &self.pid)
+			.field("exe", &self.exe)
+			.field("dispatch_join_handle", &self.dispatch_join_handle)
+			.field("flush_join_handle", &self.flush_join_handle)
+			.field("disconnect_status", &self.disconnect_status)
+			.field("base_resource_prefixes", &self.base_resource_prefixes)
+			.field("state", &self.state)
+			.finish()
+	}
+}
 impl Drop for Client {
 	fn drop(&mut self) {
 		info!(
--- a/src/core/client_state.rs
+++ b/src/core/client_state.rs
@@ -4,7 +4,11 @@ use glam::Mat4;
 use parking_lot::Mutex;
 use rustc_hash::FxHashMap;
 use serde::{Deserialize, Serialize};
-use std::{io::Write, path::PathBuf, sync::Arc};
+use std::{
+	path::{Path, PathBuf},
+	process::Command,
+	sync::Arc,
+};

 lazy_static::lazy_static! {
 	pub static ref CLIENT_STATES: Mutex<FxHashMap<String, Arc<ClientState>>> = Default::default();
@@ -58,14 +62,23 @@ impl ClientState {
 		CLIENT_STATES.lock().insert(token.clone(), Arc::new(self));
 		token
 	}
-	pub fn to_file(self) {
-		let project_dirs = directories::ProjectDirs::from("", "", "stardust").unwrap();
-		let state_dir = project_dirs.state_dir().unwrap();
-		std::fs::create_dir_all(state_dir).unwrap();
-		let mut file = std::fs::File::create(state_dir.join(nanoid::nanoid!())).unwrap();
-		file.write_all(&stardust_xr::schemas::flex::flexbuffers::to_vec(self).unwrap())
-			.unwrap();
+	pub fn from_file(file: &Path) -> Option<Self> {
+		let file_string = std::fs::read_to_string(file).ok()?;
+		toml::from_str(&file_string).ok()
 	}
+	pub fn to_file(self, directory: &Path) {
+		let app_name = self
+			.launch_info
+			.as_ref()
+			.map(|l| l.cmdline.get(0).unwrap().split('/').last().unwrap())
+			.unwrap_or("unknown");
+		let state_file_path = directory
+			.join(format!("{app_name}-{}", nanoid::nanoid!()))
+			.with_extension("toml");
+
+		std::fs::write(state_file_path, toml::to_string(&self).unwrap()).unwrap();
+	}
+
 	pub fn apply_to(&self, client: &Arc<Client>) -> ClientStateInternal {
 		if let Some(root) = client.root.get() {
 			root.set_transform(self.root)
@@ -88,6 +101,16 @@ impl ClientState {
 				.collect(),
 		}
 	}
+	pub fn launch_command(self) -> Option<Command> {
+		let launch_info = self.launch_info.as_ref()?;
+		let mut cmdline = launch_info.cmdline.iter();
+		let mut command = Command::new(cmdline.next()?);
+		command.args(cmdline);
+		command.current_dir(&launch_info.cwd);
+		command.envs(launch_info.env.iter());
+		command.env("STARDUST_STARTUP_TOKEN", self.token());
+		Some(command)
+	}
 }
 impl Default for ClientState {
 	fn default() -> Self {
--- a/src/main.rs
+++ b/src/main.rs
@@ -5,7 +5,6 @@ mod objects;
 mod wayland;

 use crate::core::client::CLIENTS;
-use crate::core::client_state::ClientState;
 use crate::core::destroy_queue;
 use crate::nodes::items::camera;
 use crate::nodes::{audio, drawable, hmd, input};
@@ -18,12 +17,13 @@ use crate::wayland::X_DISPLAY;

 use self::core::eventloop::EventLoop;
 use clap::Parser;
+use core::client_state::ClientState;
 use directories::ProjectDirs;
 use once_cell::sync::OnceCell;
 use stardust_xr::server;
-use std::os::unix::process::CommandExt;
-use std::path::PathBuf;
-use std::process::{Command, Stdio};
+use std::os::unix::fs::PermissionsExt;
+use std::path::{Path, PathBuf};
+use std::process::{Child, Command, Stdio};
 use std::sync::Arc;
 use std::time::Duration;
 use stereokit::{
@@ -56,6 +56,10 @@ struct CliArgs {
 	/// Run a script when ready for clients to connect. If this is not set the script at $HOME/.config/stardust/startup will be ran if it exists.
 	#[clap(id = "PATH", short = 'e', long = "execute-startup-script", action)]
 	startup_script: Option<PathBuf>,
+
+	/// Restore the session with the given ID (or `latest`), ignoring the startup script. Sessions are stored in directories at `~/.local/state/stardust/`.
+	#[clap(id = "SESSION_ID", long = "restore", action)]
+	restore: Option<String>,
 }

 static STARDUST_INSTANCE: OnceCell<String> = OnceCell::new();
@@ -197,7 +201,10 @@ fn main() {
 	let (info_sender, info_receiver) = oneshot::channel::<EventLoopInfo>();
 	let event_thread = std::thread::Builder::new()
 		.name("event_loop".to_owned())
-		.spawn(move || event_loop(info_sender))
+		.spawn({
+			let project_dirs = project_dirs.clone();
+			move || event_loop(info_sender, project_dirs.clone())
+		})
 		.unwrap();
 	let event_loop_info = info_receiver.blocking_recv().unwrap();
 	let _tokio_handle = event_loop_info.tokio_handle.enter();
@@ -206,56 +213,18 @@ fn main() {
 	let mut wayland = wayland::Wayland::new().expect("Could not initialize wayland");
 	info!("Stardust ready!");

-	let mut startup_child = (|| {
-		let project_dirs = project_dirs.as_ref()?;
-		let startup_script_path = cli_args
-			.startup_script
-			.clone()
-			.and_then(|p| p.canonicalize().ok())
-			.unwrap_or_else(|| project_dirs.config_dir().join("startup"));
-		let mut startup_command = Command::new("bash");
-		startup_command.arg(startup_script_path);
-		startup_command.arg("&");
-
-		startup_command.stdin(Stdio::null());
-		startup_command.stdout(Stdio::null());
-		startup_command.stderr(Stdio::null());
-		startup_command.env(
-			"FLAT_WAYLAND_DISPLAY",
-			std::env::var_os("WAYLAND_DISPLAY").unwrap_or_default(),
-		);
-		startup_command.env(
-			"STARDUST_INSTANCE",
-			event_loop_info
-				.socket_path
-				.file_name()
-				.expect("Stardust socket path not found"),
-		);
-		#[cfg(feature = "wayland")]
-		{
-			if let Some(wayland_socket) = wayland.socket_name.as_ref() {
-				startup_command.env("WAYLAND_DISPLAY", &wayland_socket);
-			}
-			startup_command.env(
-				"DISPLAY",
-				format!(":{}", X_DISPLAY.get().cloned().unwrap_or_default()),
-			);
-			startup_command.env("GDK_BACKEND", "wayland");
-			startup_command.env("QT_QPA_PLATFORM", "wayland");
-			startup_command.env("MOZ_ENABLE_WAYLAND", "1");
-			startup_command.env("CLUTTER_BACKEND", "wayland");
-			startup_command.env("SDL_VIDEODRIVER", "wayland");
-		}
-		unsafe {
-			startup_command.pre_exec(|| {
-				nix::unistd::setsid()
-					.map(|_| ())
-					.map_err(|_| std::io::ErrorKind::Other.into())
-			})
-		};
-		let child = startup_command.spawn().ok()?;
-		Some(child)
-	})();
+	let mut startup_children = project_dirs
+		.as_ref()
+		.map(|project_dirs| {
+			launch_start(
+				&cli_args,
+				project_dirs,
+				&event_loop_info,
+				#[cfg(feature = "wayland")]
+				&wayland,
+			)
+		})
+		.unwrap_or_default();

 	let mut last_frame_delta = Duration::ZERO;
 	let mut sleep_duration = Duration::ZERO;
@@ -306,10 +275,6 @@ fn main() {
 			},
 			|_sk| {
 				info!("Cleanly shut down StereoKit");
-
-				if let Some(mut startup_child) = startup_child.take() {
-					let _ = startup_child.kill();
-				}
 			},
 		)
 	});
@@ -322,6 +287,9 @@ fn main() {
 		.join()
 		.expect("Failed to cleanly shut down event loop")
 		.unwrap();
+	for mut startup_child in startup_children.drain(..) {
+		let _ = startup_child.kill();
+	}

 	info!("Cleanly shut down Stardust");
 }
@@ -351,7 +319,10 @@ fn adaptive_sleep(

 // #[tokio::main]
 #[tokio::main(flavor = "current_thread")]
-async fn event_loop(info_sender: oneshot::Sender<EventLoopInfo>) -> color_eyre::eyre::Result<()> {
+async fn event_loop(
+	info_sender: oneshot::Sender<EventLoopInfo>,
+	project_dirs: Option<ProjectDirs>,
+) -> color_eyre::eyre::Result<()> {
 	let socket_path =
 		server::get_free_socket_path().expect("Unable to find a free stardust socket path");
 	STARDUST_INSTANCE.set(socket_path.file_name().unwrap().to_string_lossy().into_owned()).expect("Someone hasn't done their job, yell at Nova because how is this set multiple times what the hell");
@@ -368,7 +339,9 @@ async fn event_loop(info_sender: oneshot::Sender<EventLoopInfo>) -> color_eyre::

 	STOP_NOTIFIER.notified().await;
 	println!("Stopping...");
-	save_clients().await;
+	if let Some(project_dirs) = project_dirs {
+		save_session(&project_dirs).await;
+	}

 	info!("Cleanly shut down event loop");

@@ -379,16 +352,121 @@ async fn event_loop(info_sender: oneshot::Sender<EventLoopInfo>) -> color_eyre::
 	Ok(())
 }

-async fn save_clients() {
+fn launch_start(
+	cli_args: &CliArgs,
+	project_dirs: &ProjectDirs,
+	event_loop_info: &EventLoopInfo,
+	#[cfg(feature = "wayland")] wayland: &wayland::Wayland,
+) -> Vec<Child> {
+	if let Some(session_id) = &cli_args.restore {
+		let session_dir = project_dirs.state_dir().unwrap().join(session_id);
+		return restore_session(&session_dir, event_loop_info, wayland);
+	}
+	let startup_script_path = cli_args
+		.startup_script
+		.clone()
+		.and_then(|p| p.canonicalize().ok())
+		.unwrap_or_else(|| project_dirs.config_dir().join("startup"));
+	run_script(&startup_script_path, event_loop_info, wayland)
+}
+
+fn restore_session(
+	session_dir: &Path,
+	event_loop_info: &EventLoopInfo,
+	#[cfg(feature = "wayland")] wayland: &wayland::Wayland,
+) -> Vec<Child> {
+	let Ok(clients) = session_dir.read_dir() else {
+		return Vec::new();
+	};
+	clients
+		.filter_map(Result::ok)
+		.filter_map(|c| ClientState::from_file(&c.path()))
+		.filter_map(ClientState::launch_command)
+		.filter_map(|startup_command| {
+			run_client(
+				startup_command,
+				event_loop_info,
+				#[cfg(feature = "wayland")]
+				wayland,
+			)
+		})
+		.collect()
+}
+
+fn run_script(
+	script_path: &Path,
+	event_loop_info: &EventLoopInfo,
+	#[cfg(feature = "wayland")] wayland: &wayland::Wayland,
+) -> Vec<Child> {
+	let _ = std::fs::set_permissions(script_path, std::fs::Permissions::from_mode(0o755));
+	let startup_command = Command::new(script_path);
+	run_client(
+		startup_command,
+		event_loop_info,
+		#[cfg(feature = "wayland")]
+		wayland,
+	)
+	.map(|c| vec![c])
+	.unwrap_or_default()
+}
+
+fn run_client(
+	mut command: Command,
+	event_loop_info: &EventLoopInfo,
+	#[cfg(feature = "wayland")] wayland: &wayland::Wayland,
+) -> Option<Child> {
+	command.stdin(Stdio::null());
+	command.stdout(Stdio::null());
+	command.stderr(Stdio::null());
+	command.env(
+		"FLAT_WAYLAND_DISPLAY",
+		std::env::var_os("WAYLAND_DISPLAY").unwrap_or_default(),
+	);
+	command.env(
+		"STARDUST_INSTANCE",
+		event_loop_info
+			.socket_path
+			.file_name()
+			.expect("Stardust socket path not found"),
+	);
+	#[cfg(feature = "wayland")]
+	{
+		if let Some(wayland_socket) = wayland.socket_name.as_ref() {
+			command.env("WAYLAND_DISPLAY", &wayland_socket);
+		}
+		command.env(
+			"DISPLAY",
+			format!(":{}", X_DISPLAY.get().cloned().unwrap_or_default()),
+		);
+		command.env("GDK_BACKEND", "wayland");
+		command.env("QT_QPA_PLATFORM", "wayland");
+		command.env("MOZ_ENABLE_WAYLAND", "1");
+		command.env("CLUTTER_BACKEND", "wayland");
+		command.env("SDL_VIDEODRIVER", "wayland");
+	}
+	let child = command.spawn().ok()?;
+	Some(child)
+}
+
+async fn save_session(project_dirs: &ProjectDirs) {
+	let session_id = nanoid::nanoid!();
+	let state_dir = project_dirs.state_dir().unwrap();
+	let session_dir = state_dir.join(&session_id);
+	std::fs::create_dir_all(&session_dir).unwrap();
+	let _ = std::fs::remove_dir_all(state_dir.join("latest"));
+	std::os::unix::fs::symlink(&session_dir, state_dir.join("latest")).unwrap();
+
 	let local_set = LocalSet::new();
 	for client in CLIENTS.get_vec() {
+		let session_dir = session_dir.clone();
 		local_set.spawn_local(async move {
 			tokio::select! {
 				biased;
-				s = client.save_state() => {s.map(ClientState::to_file);},
+				s = client.save_state() => {s.map(|s| s.to_file(&session_dir));},
 				_ = tokio::time::sleep(Duration::from_millis(100)) => (),
 			}
 		});
 	}
 	local_set.await;
+	println!("Session ID for restore is {session_id}");
 }
--- a/src/nodes/root.rs
+++ b/src/nodes/root.rs
@@ -95,11 +95,10 @@ impl Root {
 		spatial.set_local_transform(transform);
 	}
 	pub async fn save_state(&self) -> Result<ClientStateInternal> {
-		Ok(self
-			.node
+		self.node
 			.execute_remote_method_typed("save_state", (), Vec::new())
-			.await?
-			.0)
+			.await
+			.map(|(m, _)| m)
 	}
 }