[PATCH v1 2/2] rust: task: add Rust version of might_sleep()

[FUJITA Tomonori]

Adds a helper function equivalent to the C's might_sleep(), which
serves as a debugging aid and a potential scheduling point.

Note that this function can only be used in a nonatomic context.

This will be used by Rust version of read_poll_timeout().

Signed-off-by: FUJITA Tomonori <fujita.tomonori@xxxxxxxxx>
---
 rust/helpers/task.c |  6 ++++++
 rust/kernel/task.rs | 26 ++++++++++++++++++++++++++
 2 files changed, 32 insertions(+)

diff --git a/rust/helpers/task.c b/rust/helpers/task.c
index 31c33ea2dce6..2c85bbc2727e 100644
--- a/rust/helpers/task.c
+++ b/rust/helpers/task.c
@@ -1,7 +1,13 @@
 // SPDX-License-Identifier: GPL-2.0
 
+#include <linux/kernel.h>
 #include <linux/sched/task.h>
 
+void rust_helper_might_resched(void)
+{
+	might_resched();
+}
+
 struct task_struct *rust_helper_get_current(void)
 {
 	return current;
diff --git a/rust/kernel/task.rs b/rust/kernel/task.rs
index 9e6f6854948d..1f0156b38ab5 100644
--- a/rust/kernel/task.rs
+++ b/rust/kernel/task.rs
@@ -380,3 +380,29 @@ fn eq(&self, other: &Kuid) -> bool {
 }
 
 impl Eq for Kuid {}
+
+/// Annotation for functions that can sleep.
+///
+/// Equivalent to the C side [`might_sleep()`], this function serves as
+/// a debugging aid and a potential scheduling point.
+///
+/// This function can only be used in a nonatomic context.
+#[track_caller]
+#[inline]
+pub fn might_sleep() {
+    #[cfg(CONFIG_DEBUG_ATOMIC_SLEEP)]
+    {
+        let loc = core::panic::Location::caller();
+        // SAFETY: FFI call.
+        unsafe {
+            crate::bindings::__might_sleep_precision(
+                loc.file().as_ptr().cast(),
+                loc.file().len() as i32,
+                loc.line() as i32,
+            )
+        }
+    }
+
+    // SAFETY: FFI call.
+    unsafe { crate::bindings::might_resched() }
+}
-- 
2.43.0