fix(tls): create TLS for idle task when instrumenting the kernel

Cargo.toml

@@ -44,6 +44,7 @@ fs = ["virtio-fs"]
 fsgsbase = []
 gem-net = ["net", "dep:tock-registers"]
 idle-poll = []
+instrument-mcount = [] # Inserts function instrument code for mcount-based tracing
 kernel-stack = []
 net = []
 mman = []

src/scheduler/task/mod.rs

@@ -477,6 +477,13 @@ impl Task {
 			}
 		}
 
+		#[cfg(not(feature = "common-os"))]
+		let tls = if cfg!(feature = "instrument-mcount") {
+			Tls::from_env().inspect(Tls::set_thread_ptr)
+		} else {
+			None
+		};
+
 		Task {
 			id: tid,
 			status: TaskStatus::Idle,
@@ -488,7 +495,7 @@ impl Task {
 			stacks: TaskStacks::from_boot_stacks(),
 			object_map: OBJECT_MAP.get().unwrap().clone(),
 			#[cfg(not(feature = "common-os"))]
-			tls: None,
+			tls,
 			#[cfg(all(target_arch = "x86_64", feature = "common-os"))]
 			root_page_table: *crate::scheduler::BOOT_ROOT_PAGE_TABLE.get().unwrap(),
 		}

src/scheduler/task/tls.rs

@@ -146,6 +146,41 @@ impl Tls {
 	pub fn thread_ptr(&self) -> *mut () {
 		self.thread_ptr
 	}
+
+	/// Sets the thread pointer register to this TLS's thread pointer value.
+	///
+	/// This should only be used for the idle task. Since the idle task is
+	/// already running, we don't create a new stack frame that we would put
+	/// the thread pointer value into otherwise.
+	///
+	/// The idle task does not enter userspace. That's why it does not need a
+	/// TLS most of the time. In special situations such as instrumenting the
+	/// kernel, the tracer or profiler or `mcount` implementation such as
+	/// rftrace might use TLS for differentiating between the idle task and
+	/// other tasks.
+	pub fn set_thread_ptr(&self) {
+		cfg_if::cfg_if! {
+			if #[cfg(target_arch = "aarch64")] {
+				use aarch64_cpu::registers::{TPIDR_EL0, Writeable};
+
+				let addr = self.thread_ptr().expose_provenance();
+				TPIDR_EL0.set(addr.try_into().unwrap());
+			} else if #[cfg(target_arch = "riscv64")] {
+				unsafe {
+					core::arch::asm!(
+						"mv tp, {}",
+						in(reg) self.thread_ptr().expose_provenance(),
+						options(nomem, nostack, preserves_flags),
+					);
+				}
+			} else if #[cfg(target_arch = "x86_64")] {
+				use crate::arch::x86_64::kernel::processor;
+
+				let addr = self.thread_ptr().expose_provenance();
+				processor::writefs(addr);
+			}
+		}
+	}
 }
 
 mod allocation {

xtask/src/build.rs

@@ -12,7 +12,7 @@ pub struct Build {
 	#[command(flatten)]
 	cargo_build: CargoBuild,
 
-	/// Enable the `-Z instrument-mcount` flag.
+	/// Deprecated: use `--features instrument-mcount` instead.
 	#[arg(long)]
 	pub instrument_mcount: bool,
 
@@ -22,9 +22,15 @@ pub struct Build {
 }
 
 impl Build {
-	pub fn run(self) -> Result<()> {
+	pub fn run(mut self) -> Result<()> {
 		let sh = crate::sh()?;
 
+		if self.instrument_mcount {
+			self.cargo_build
+				.features
+				.push("instrument-mcount".to_owned());
+		}
+
 		self.cargo_build.artifact.arch.install_for_build()?;
 
 		let careful = match env::var_os("HERMIT_CAREFUL") {
@@ -96,7 +102,12 @@ impl Build {
 			.map(|s| vec![s])
 			.unwrap_or_default();
 
-		if self.instrument_mcount {
+		if self
+			.cargo_build
+			.features
+			.iter()
+			.any(|feature| feature == "instrument-mcount")
+		{
 			rustflags.push("-Zinstrument-mcount");
 			rustflags.push("-Cpasses=ee-instrument<post-inline>");
 		}