..

Author: sxyazi

yazi-plugin/src/process/child.rs

@@ -9,18 +9,39 @@ use super::Status;
 use crate::process::Output;
 
 pub struct Child {
-	inner:  tokio::process::Child,
-	stdin:  Option<BufWriter<ChildStdin>>,
-	stdout: Option<BufReader<ChildStdout>>,
-	stderr: Option<BufReader<ChildStderr>>,
+	inner:      tokio::process::Child,
+	stdin:      Option<BufWriter<ChildStdin>>,
+	stdout:     Option<BufReader<ChildStdout>>,
+	stderr:     Option<BufReader<ChildStderr>>,
+	#[cfg(windows)]
+	job_handle: Option<std::os::windows::io::RawHandle>,
+}
+
+#[cfg(windows)]
+impl Drop for Child {
+	fn drop(&mut self) {
+		if let Some(h) = self.job_handle.take() {
+			unsafe { windows_sys::Win32::Foundation::CloseHandle(h) };
+		}
+	}
 }
 
 impl Child {
-	pub fn new(mut inner: tokio::process::Child) -> Self {
+	pub fn new(
+		mut inner: tokio::process::Child,
+		#[cfg(windows)] job_handle: Option<std::os::windows::io::RawHandle>,
+	) -> Self {
 		let stdin = inner.stdin.take().map(BufWriter::new);
 		let stdout = inner.stdout.take().map(BufReader::new);
 		let stderr = inner.stderr.take().map(BufReader::new);
-		Self { inner, stdin, stdout, stderr }
+		Self {
+			inner,
+			stdin,
+			stdout,
+			stderr,
+			#[cfg(windows)]
+			job_handle,
+		}
 	}
 
 	pub(super) async fn wait(&mut self) -> io::Result<ExitStatus> {

yazi-plugin/src/process/command.rs

@@ -43,9 +43,7 @@ impl Command {
 			unsafe {
 				self.inner.pre_exec(move || {
 					let rlp = libc::rlimit { rlim_cur: max as _, rlim_max: max as _ };
-					if libc::setrlimit(libc::RLIMIT_AS, &rlp) != 0 {
-						return Err(io::Error::last_os_error());
-					}
+					libc::setrlimit(libc::RLIMIT_AS, &rlp);
 					Ok(())
 				});
 			}
@@ -55,29 +53,46 @@ impl Command {
 
 	#[cfg(windows)]
 	fn spawn(&mut self) -> io::Result<Child> {
-		unsafe {
-			let job = CreateJobObjectW(std::ptr::null_mut(), std::ptr::null());
-
-			let mut info: JOBOBJECT_EXTENDED_LIMIT_INFORMATION = std::mem::zeroed();
-			info.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_PROCESS_MEMORY;
-			info.ProcessMemoryLimit = 536870912; // 512MB
-
-			SetInformationJobObject(
-				job,
-				JobObjectExtendedLimitInformation,
-				&mut info as *mut _ as *mut _,
-				std::mem::size_of::<JOBOBJECT_EXTENDED_LIMIT_INFORMATION>() as u32,
-			);
+		use std::os::windows::io::RawHandle;
 
-			let mut child =
-				Command::new("your_program").args(&["arg1", "arg2"]).stdout(Stdio::piped()).spawn()?;
+		use windows_sys::Win32::System::JobObjects::{AssignProcessToJobObject, CreateJobObjectW, JOB_OBJECT_LIMIT_PROCESS_MEMORY, JOBOBJECT_EXTENDED_LIMIT_INFORMATION, JobObjectExtendedLimitInformation, SetInformationJobObject};
 
-			AssignProcessToJobObject(job, child.as_raw_handle() as HANDLE);
+		fn assign_job(max: usize, handle: RawHandle) -> io::Result<RawHandle> {
+			unsafe {
+				let job = CreateJobObjectW(std::ptr::null_mut(), std::ptr::null());
+				if job.is_null() {
+					return Err(io::Error::last_os_error());
+				}
 
-			let status = child.wait()?;
+				let mut info: JOBOBJECT_EXTENDED_LIMIT_INFORMATION = std::mem::zeroed();
+				info.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_PROCESS_MEMORY;
+				info.ProcessMemoryLimit = max;
+
+				let result = SetInformationJobObject(
+					job,
+					JobObjectExtendedLimitInformation,
+					&mut info as *mut _ as *mut _,
+					std::mem::size_of_val(&info) as u32,
+				);
+
+				if result == 0 {
+					Err(io::Error::last_os_error())
+				} else if AssignProcessToJobObject(job, handle) == 0 {
+					Err(io::Error::last_os_error())
+				} else {
+					Ok(job)
+				}
+			}
+		}
 
-			CloseHandle(job);
+		let child = self.inner.spawn()?;
+		if let (Some(max), Some(handle)) = (self.memory, child.raw_handle()) {
+			if let Ok(job) = assign_job(max, handle) {
+				return Ok(Child::new(child, Some(job)));
+			}
 		}
+
+		Ok(Child::new(child, None))
 	}
 
 	async fn output(&mut self) -> io::Result<std::process::Output> {