resourced/schedqos/src/proc.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.

 use std::fmt::Display;
 use std::fs::File;
 use std::io;
 use std::io::BufRead;
 use std::io::BufReader;
 use std::io::Read;
 use std::path::Path;

 use crate::ProcessId;
 use crate::ThreadId;

 pub type Result<T> = std::result::Result<T, Error>;

 #[derive(Debug)]
 pub enum Error {
     Io(io::Error),
     NotFound,
     FormatCorrupt,
 }

 impl std::error::Error for Error {
     fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
         match self {
             Self::Io(e) => Some(e),
             _ => None,
         }
     }
 }

 impl Display for Error {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
             Self::Io(e) => e.fmt(f),
             Self::NotFound => f.write_str("process/thread not found"),
             Self::FormatCorrupt => f.write_str("stat file is corrupt"),
         }
     }
 }

 impl From<io::Error> for Error {
     fn from(e: io::Error) -> Self {
         if e.kind() == io::ErrorKind::NotFound || e.raw_os_error() == Some(libc::ESRCH) {
             Self::NotFound
         } else {
             Self::Io(e)
         }
     }
 }

 pub fn load_process_timestamp(process_id: ProcessId) -> Result<u64> {
     load_starttime(Path::new(&format!("/proc/{}/stat", process_id.0)))
 }

 pub fn load_thread_timestamp(process_id: ProcessId, thread_id: ThreadId) -> Result<u64> {
     load_starttime(Path::new(&format!(
         "/proc/{}/task/{}/stat",
         process_id.0, thread_id.0
     )))
 }

 fn load_starttime(path: &Path) -> Result<u64> {
     let mut stat_file = File::open(path)?;
     // starttime is the 22th column in /proc/pid/stat. Each numeric column in /proc/pid/stat has at
     // most 21 bytes. (1 byte for sign + 19 bytes for u64 + 1 byte space). The 2nd column (comm) is
     // at most 67 bytes including the wrapping parenthesis (proc_task_name() of kernel uses 64 bytes
     // buffer `tcomm`). 512 bytes is enough to hold the 22 columns (i.e. 512 >= 21 * 21 + 67 = 508).
     let mut buf = [0; 512];
     let n = stat_file.read(&mut buf)?;

     // Since threads can set comm by writing to /proc/self/task/tid/comm, it can contain any byte
     // sequence. This means we need to exclusively look at kernel controlled bytes to determine
     // where comm ends. To do this, we can scan backwards to look for the closing parentheses
     // emitted by the kernel.
     // The longest possible tail offset of comm is 88 bytes (= 21 bytes for pid column + 1 byte for
     // space + 66 bytes for comm column). This works even if the stat file size is less than 88
     // bytes because the space after the stat file content are zeroed.
     let mut i_comm_tail = None;
     for (i, c) in buf[..88].iter().enumerate().rev() {
         if *c == b')' {
             i_comm_tail = Some(i);
             break;
         }
     }
     let Some(i_comm_tail) = i_comm_tail else {
         return Err(Error::FormatCorrupt);
     };

     let mut prev_space = i_comm_tail + 1;
     let mut starttime = None;
     // `pid` and `comm` columns are consumed.
     let mut column_idx = 2;
     for (i, c) in buf[..n].iter().enumerate().skip(prev_space + 1) {
         if *c == b' ' {
             if column_idx == 21 {
                 // starttime is at 22th column.
                 starttime = Some(prev_space + 1..i);
                 break;
             }
             prev_space = i;
             column_idx += 1;
         }
     }
     let Some(starttime) = starttime else {
         return Err(Error::FormatCorrupt);
     };
     let starttime = std::str::from_utf8(&buf[starttime]).map_err(|_| Error::FormatCorrupt)?;
     let starttime = starttime.parse().map_err(|_| Error::FormatCorrupt)?;

     Ok(starttime)
 }

 pub fn load_tgid(thread_id: ThreadId) -> Result<ProcessId> {
     let file = File::open(format!("/proc/{}/status", thread_id.0))?;
     let r = BufReader::with_capacity(1024, file);
     for line in r.lines() {
         let line = line.map_err(Error::Io)?;
         const TGID_TAG: &str = "Tgid:";
         if let Some(tgid) = line.strip_prefix(TGID_TAG) {
             let tgid = tgid.trim().parse().map_err(|_| Error::FormatCorrupt)?;
             return Ok(ProcessId(tgid));
         }
     }
     Err(Error::FormatCorrupt)
 }

 pub struct ThreadChecker {
     prefix: String,
 }

 impl ThreadChecker {
     pub fn new(process_id: ProcessId) -> Self {
         // "/proc/" (6 bytes) + pid (at most 10 bytes) "/task/" (6 bytes) + tid (at most 10 bytes).
         let mut prefix = String::with_capacity(32);
         prefix.push_str("/proc/");
         prefix.push_str(&process_id.0.to_string());
         prefix.push_str("/task/");
         Self { prefix }
     }

     pub fn thread_exists(&mut self, thread_id: ThreadId) -> bool {
         let original_len = self.prefix.len();
         self.prefix.push_str(&thread_id.0.to_string());
         let result = Path::new(&self.prefix).exists();
         self.prefix.truncate(original_len);
         result
     }
 }

 #[cfg(test)]
 mod tests {
     use std::io::Write;

     use super::*;
     use crate::test_utils::*;

     #[test]
     fn test_load_process_timestamp() {
         let process_id = ProcessId(std::process::id());
         assert!(load_process_timestamp(process_id).unwrap() > 0);

         let (process_id, _, _process) = fork_process_for_test();
         assert!(load_process_timestamp(process_id).unwrap() > 0);
     }

     #[test]
     fn test_load_process_timestamp_not_found() {
         let (process_id, _, process) = fork_process_for_test();
         drop(process);
         assert!(matches!(
             load_process_timestamp(process_id).err().unwrap(),
             Error::NotFound
         ));
     }

     #[test]
     fn test_load_thread_timestamp() {
         let process_id = ProcessId(std::process::id());
         let thread_id = get_current_thread_id();
         assert!(load_thread_timestamp(process_id, thread_id).unwrap() > 0);

         let (thread_id, _thread) = spawn_thread_for_test();
         assert!(load_thread_timestamp(process_id, thread_id).unwrap() > 0);

         let (process_id, thread_id, _process) = fork_process_for_test();
         assert!(load_thread_timestamp(process_id, thread_id).unwrap() > 0);
     }

     #[test]
     fn test_load_thread_timestamp_not_found() {
         let process_id = ProcessId(std::process::id());
         let thread_id = get_current_thread_id();
         let (child_process_id, child_thread_id, _process) = fork_process_for_test();

         assert!(matches!(
             load_thread_timestamp(process_id, child_thread_id)
                 .err()
                 .unwrap(),
             Error::NotFound
         ));
         assert!(matches!(
             load_thread_timestamp(child_process_id, thread_id)
                 .err()
                 .unwrap(),
             Error::NotFound
         ));
     }

     #[test]
     fn test_load_starttime() {
         for (stat_file_content, starttime) in [
             (
                 "9345 (resourced) S 1 9344 9344 0 -1 1077936384 599 0 0 0 2851 2468 0 0 20 0 1 0 \
             41329081 19865600 2719 18446744073709551615 101386084483072 101386086716560 \
             140736188509360 0 0 0 0 4096 1088 1 0 0 17 0 0 0 0 0 0 101386086981672 101386086982232 \
             101386110365696 140736188513963 140736188513982 140736188513982 140736188514277 0",
                 41329081,
             ),
             // Shortest stat file.
             ("1 (a) 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 ", 2),
             // Longest pid + comm stat file.
             (
                 "1234567890123456789 \
             (1234567890123456789012345678901234567890123456789012345678901234) 3 4 5 6 7 8 9 0 1 2 \
             3 4 5 6 7 8 9 0 1 123 ",
                 123,
             ),
             // comm contains spaces and parenthesis.
             (
                 "1 ( a ( b ) c ) ) 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 456 ",
                 456,
             ),
         ] {
             let mut file = tempfile::NamedTempFile::new().unwrap();
             file.write_all(stat_file_content.as_bytes()).unwrap();
             assert_eq!(
                 load_starttime(file.path()).unwrap(),
                 starttime,
                 "{}",
                 stat_file_content
             );
         }
     }

     #[test]
     fn test_load_starttime_corrupt_format() {
         for stat_file_content in [
             // no comm parenthesis.
             "1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2",
             // starttime is not a number.
             "1 (a) 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 a ",
             // no 22th space
             "1 (a) 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2",
         ] {
             let mut file = tempfile::NamedTempFile::new().unwrap();
             file.write_all(stat_file_content.as_bytes()).unwrap();
             assert!(
                 matches!(
                     load_starttime(file.path()).err().unwrap(),
                     Error::FormatCorrupt
                 ),
                 "{}",
                 stat_file_content
             );
         }

         let mut file = tempfile::NamedTempFile::new().unwrap();
         file.write_all(b"1 (a) 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2")
             .unwrap();
         // non utf-8
         file.write_all(&[255, b' ']).unwrap();
         assert!(matches!(
             load_starttime(file.path()).err().unwrap(),
             Error::FormatCorrupt
         ));
     }

     #[test]
     fn test_load_tgid() {
         let process_id = ProcessId(std::process::id());
         assert_eq!(load_tgid(ThreadId(process_id.0)).unwrap(), process_id);

         let (thread_id, thread) = spawn_thread_for_test();
         assert_eq!(load_tgid(thread_id).unwrap(), process_id);
         drop(thread);
         wait_for_thread_removed(process_id, thread_id);
         assert!(matches!(load_tgid(thread_id), Err(Error::NotFound)));

         let (process_id, thread_id, _process) = fork_process_for_test();
         assert_eq!(load_tgid(thread_id).unwrap(), process_id);
     }

     #[test]
     fn test_thread_exists() {
         let process_id = ProcessId(std::process::id());
         let mut checker = ThreadChecker::new(process_id);
         let (thread_id, thread) = spawn_thread_for_test();
         let (another_process_id, another_thread_id, process) = fork_process_for_test();
         let mut another_checker = ThreadChecker::new(another_process_id);

         assert!(checker.thread_exists(thread_id));
         assert!(!checker.thread_exists(another_thread_id));
         assert!(another_checker.thread_exists(another_thread_id));
         assert!(!another_checker.thread_exists(thread_id));
         drop(thread);
         wait_for_thread_removed(process_id, thread_id);
         assert!(!checker.thread_exists(thread_id));
         drop(process);
         assert!(!another_checker.thread_exists(another_thread_id));
     }
 }
	// 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.

	use std::fmt::Display;
	use std::fs::File;
	use std::io;
	use std::io::BufRead;
	use std::io::BufReader;
	use std::io::Read;
	use std::path::Path;

	use crate::ProcessId;
	use crate::ThreadId;

	pub type Result<T> = std::result::Result<T, Error>;

	#[derive(Debug)]
	pub enum Error {
	Io(io::Error),
	NotFound,
	FormatCorrupt,
	}

	impl std::error::Error for Error {
	fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
	match self {
	Self::Io(e) => Some(e),
	_ => None,
	}
	}
	}

	impl Display for Error {
	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
	match self {
	Self::Io(e) => e.fmt(f),
	Self::NotFound => f.write_str("process/thread not found"),
	Self::FormatCorrupt => f.write_str("stat file is corrupt"),
	}
	}
	}

	impl From<io::Error> for Error {
	fn from(e: io::Error) -> Self {
	if e.kind() == io::ErrorKind::NotFound \|\| e.raw_os_error() == Some(libc::ESRCH) {
	Self::NotFound
	} else {
	Self::Io(e)
	}
	}
	}

	pub fn load_process_timestamp(process_id: ProcessId) -> Result<u64> {
	load_starttime(Path::new(&format!("/proc/{}/stat", process_id.0)))
	}

	pub fn load_thread_timestamp(process_id: ProcessId, thread_id: ThreadId) -> Result<u64> {
	load_starttime(Path::new(&format!(
	"/proc/{}/task/{}/stat",
	process_id.0, thread_id.0
	)))
	}

	fn load_starttime(path: &Path) -> Result<u64> {
	let mut stat_file = File::open(path)?;
	// starttime is the 22th column in /proc/pid/stat. Each numeric column in /proc/pid/stat has at
	// most 21 bytes. (1 byte for sign + 19 bytes for u64 + 1 byte space). The 2nd column (comm) is
	// at most 67 bytes including the wrapping parenthesis (proc_task_name() of kernel uses 64 bytes
	// buffer `tcomm`). 512 bytes is enough to hold the 22 columns (i.e. 512 >= 21 * 21 + 67 = 508).
	let mut buf = [0; 512];
	let n = stat_file.read(&mut buf)?;

	// Since threads can set comm by writing to /proc/self/task/tid/comm, it can contain any byte
	// sequence. This means we need to exclusively look at kernel controlled bytes to determine
	// where comm ends. To do this, we can scan backwards to look for the closing parentheses
	// emitted by the kernel.
	// The longest possible tail offset of comm is 88 bytes (= 21 bytes for pid column + 1 byte for
	// space + 66 bytes for comm column). This works even if the stat file size is less than 88
	// bytes because the space after the stat file content are zeroed.
	let mut i_comm_tail = None;
	for (i, c) in buf[..88].iter().enumerate().rev() {
	if *c == b')' {
	i_comm_tail = Some(i);
	break;
	}
	}
	let Some(i_comm_tail) = i_comm_tail else {
	return Err(Error::FormatCorrupt);
	};

	let mut prev_space = i_comm_tail + 1;
	let mut starttime = None;
	// `pid` and `comm` columns are consumed.
	let mut column_idx = 2;
	for (i, c) in buf[..n].iter().enumerate().skip(prev_space + 1) {
	if *c == b' ' {
	if column_idx == 21 {
	// starttime is at 22th column.
	starttime = Some(prev_space + 1..i);
	break;
	}
	prev_space = i;
	column_idx += 1;
	}
	}
	let Some(starttime) = starttime else {
	return Err(Error::FormatCorrupt);
	};
	let starttime = std::str::from_utf8(&buf[starttime]).map_err(\|_\| Error::FormatCorrupt)?;
	let starttime = starttime.parse().map_err(\|_\| Error::FormatCorrupt)?;

	Ok(starttime)
	}

	pub fn load_tgid(thread_id: ThreadId) -> Result<ProcessId> {
	let file = File::open(format!("/proc/{}/status", thread_id.0))?;
	let r = BufReader::with_capacity(1024, file);
	for line in r.lines() {
	let line = line.map_err(Error::Io)?;
	const TGID_TAG: &str = "Tgid:";
	if let Some(tgid) = line.strip_prefix(TGID_TAG) {
	let tgid = tgid.trim().parse().map_err(\|_\| Error::FormatCorrupt)?;
	return Ok(ProcessId(tgid));
	}
	}
	Err(Error::FormatCorrupt)
	}

	pub struct ThreadChecker {
	prefix: String,
	}

	impl ThreadChecker {
	pub fn new(process_id: ProcessId) -> Self {
	// "/proc/" (6 bytes) + pid (at most 10 bytes) "/task/" (6 bytes) + tid (at most 10 bytes).
	let mut prefix = String::with_capacity(32);
	prefix.push_str("/proc/");
	prefix.push_str(&process_id.0.to_string());
	prefix.push_str("/task/");
	Self { prefix }
	}

	pub fn thread_exists(&mut self, thread_id: ThreadId) -> bool {
	let original_len = self.prefix.len();
	self.prefix.push_str(&thread_id.0.to_string());
	let result = Path::new(&self.prefix).exists();
	self.prefix.truncate(original_len);
	result
	}
	}

	#[cfg(test)]
	mod tests {
	use std::io::Write;

	use super::*;
	use crate::test_utils::*;

	#[test]
	fn test_load_process_timestamp() {
	let process_id = ProcessId(std::process::id());
	assert!(load_process_timestamp(process_id).unwrap() > 0);

	let (process_id, _, _process) = fork_process_for_test();
	assert!(load_process_timestamp(process_id).unwrap() > 0);
	}

	#[test]
	fn test_load_process_timestamp_not_found() {
	let (process_id, _, process) = fork_process_for_test();
	drop(process);
	assert!(matches!(
	load_process_timestamp(process_id).err().unwrap(),
	Error::NotFound
	));
	}

	#[test]
	fn test_load_thread_timestamp() {
	let process_id = ProcessId(std::process::id());
	let thread_id = get_current_thread_id();
	assert!(load_thread_timestamp(process_id, thread_id).unwrap() > 0);

	let (thread_id, _thread) = spawn_thread_for_test();
	assert!(load_thread_timestamp(process_id, thread_id).unwrap() > 0);

	let (process_id, thread_id, _process) = fork_process_for_test();
	assert!(load_thread_timestamp(process_id, thread_id).unwrap() > 0);
	}

	#[test]
	fn test_load_thread_timestamp_not_found() {
	let process_id = ProcessId(std::process::id());
	let thread_id = get_current_thread_id();
	let (child_process_id, child_thread_id, _process) = fork_process_for_test();

	assert!(matches!(
	load_thread_timestamp(process_id, child_thread_id)
	.err()
	.unwrap(),
	Error::NotFound
	));
	assert!(matches!(
	load_thread_timestamp(child_process_id, thread_id)
	.err()
	.unwrap(),
	Error::NotFound
	));
	}

	#[test]
	fn test_load_starttime() {
	for (stat_file_content, starttime) in [
	(
	"9345 (resourced) S 1 9344 9344 0 -1 1077936384 599 0 0 0 2851 2468 0 0 20 0 1 0 \
	41329081 19865600 2719 18446744073709551615 101386084483072 101386086716560 \
	140736188509360 0 0 0 0 4096 1088 1 0 0 17 0 0 0 0 0 0 101386086981672 101386086982232 \
	101386110365696 140736188513963 140736188513982 140736188513982 140736188514277 0",
	41329081,
	),
	// Shortest stat file.
	("1 (a) 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 ", 2),
	// Longest pid + comm stat file.
	(
	"1234567890123456789 \
	(1234567890123456789012345678901234567890123456789012345678901234) 3 4 5 6 7 8 9 0 1 2 \
	3 4 5 6 7 8 9 0 1 123 ",
	123,
	),
	// comm contains spaces and parenthesis.
	(
	"1 ( a ( b ) c ) ) 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 456 ",
	456,
	),
	] {
	let mut file = tempfile::NamedTempFile::new().unwrap();
	file.write_all(stat_file_content.as_bytes()).unwrap();
	assert_eq!(
	load_starttime(file.path()).unwrap(),
	starttime,
	"{}",
	stat_file_content
	);
	}
	}

	#[test]
	fn test_load_starttime_corrupt_format() {
	for stat_file_content in [
	// no comm parenthesis.
	"1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2",
	// starttime is not a number.
	"1 (a) 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 a ",
	// no 22th space
	"1 (a) 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2",
	] {
	let mut file = tempfile::NamedTempFile::new().unwrap();
	file.write_all(stat_file_content.as_bytes()).unwrap();
	assert!(
	matches!(
	load_starttime(file.path()).err().unwrap(),
	Error::FormatCorrupt
	),
	"{}",
	stat_file_content
	);
	}

	let mut file = tempfile::NamedTempFile::new().unwrap();
	file.write_all(b"1 (a) 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2")
	.unwrap();
	// non utf-8
	file.write_all(&[255, b' ']).unwrap();
	assert!(matches!(
	load_starttime(file.path()).err().unwrap(),
	Error::FormatCorrupt
	));
	}

	#[test]
	fn test_load_tgid() {
	let process_id = ProcessId(std::process::id());
	assert_eq!(load_tgid(ThreadId(process_id.0)).unwrap(), process_id);

	let (thread_id, thread) = spawn_thread_for_test();
	assert_eq!(load_tgid(thread_id).unwrap(), process_id);
	drop(thread);
	wait_for_thread_removed(process_id, thread_id);
	assert!(matches!(load_tgid(thread_id), Err(Error::NotFound)));

	let (process_id, thread_id, _process) = fork_process_for_test();
	assert_eq!(load_tgid(thread_id).unwrap(), process_id);
	}

	#[test]
	fn test_thread_exists() {
	let process_id = ProcessId(std::process::id());
	let mut checker = ThreadChecker::new(process_id);
	let (thread_id, thread) = spawn_thread_for_test();
	let (another_process_id, another_thread_id, process) = fork_process_for_test();
	let mut another_checker = ThreadChecker::new(another_process_id);

	assert!(checker.thread_exists(thread_id));
	assert!(!checker.thread_exists(another_thread_id));
	assert!(another_checker.thread_exists(another_thread_id));
	assert!(!another_checker.thread_exists(thread_id));
	drop(thread);
	wait_for_thread_removed(process_id, thread_id);
	assert!(!checker.thread_exists(thread_id));
	drop(process);
	assert!(!another_checker.thread_exists(another_thread_id));
	}
	}