tests: use test kmod for more stack trace tests

Rather than the fuzzy guesses we do with a task in sigwait, take
Stephen's idea from commit 5f3a91f80d ("Add StackFrame.locals()
method") further and use the test kmod to set up some precise stack
frames that we can test. Also use kthread_park() to make sure we can't
race with the thread not being scheduled out (unlikely as that may be).

Signed-off-by: Omar Sandoval <osandov@osandov.com>

tests/linux_kernel/kmod/drgn_test.c

@@ -8,6 +8,7 @@
 // This is intended to be used with drgn's vmtest framework, but in theory it
 // can be used with any kernel that has debug info enabled (at your own risk).
 
+#include <linux/completion.h>
 #include <linux/io.h>
 #include <linux/kernel.h>
 #include <linux/kthread.h>
@@ -388,41 +389,68 @@ static int drgn_test_slab_init(void)
 
 // kthread for stack trace
 
-static struct task_struct *drgn_kthread;
+static struct task_struct *drgn_test_kthread;
+// Completion indicating that the kthread has set up its stack frames and is
+// ready to be parked.
+static DECLARE_COMPLETION(drgn_test_kthread_ready);
 
-static int __attribute__((optimize("O0"))) drgn_kthread_fn(void *arg)
+ __attribute__((__optimize__("O0")))
+static void drgn_test_kthread_fn3(void)
 {
-	int a, b, c;
+	// Create some local variables for the test cases to use. Use volatile
+	// to make doubly sure that they aren't optimized out.
+	volatile int a, b, c;
+	a = 1;
+	b = 2;
+	c = 3;
 
-retry:
-	for (a = 0; a < 100; a++) {
-		for (b = 100; b >= 0; b--) {
-			for (c = 0; c < a; c++) {
-				set_current_state(TASK_UNINTERRUPTIBLE);
-				schedule();
-				if (kthread_should_stop())
-					return 0;
-			}
+	complete(&drgn_test_kthread_ready);
+	for (;;) {
+		set_current_state(TASK_INTERRUPTIBLE);
+		if (kthread_should_stop()) {
+			__set_current_state(TASK_RUNNING);
+			break;
 		}
+		if (kthread_should_park()) {
+			__set_current_state(TASK_RUNNING);
+			kthread_parkme();
+			continue;
+		}
+		schedule();
+		__set_current_state(TASK_RUNNING);
 	}
-	goto retry;
+}
+
+ __attribute__((__optimize__("O0")))
+static void drgn_test_kthread_fn2(void)
+{
+	drgn_test_kthread_fn3();
+}
+
+ __attribute__((__optimize__("O0")))
+static int drgn_test_kthread_fn(void *arg)
+{
+	drgn_test_kthread_fn2();
+	return 0;
 }
 
 static void drgn_test_stack_trace_exit(void)
 {
-	if (drgn_kthread) {
-		kthread_stop(drgn_kthread);
-		drgn_kthread = NULL;
+	if (drgn_test_kthread) {
+		kthread_stop(drgn_test_kthread);
+		drgn_test_kthread = NULL;
 	}
 }
 
 static int drgn_test_stack_trace_init(void)
 {
-	drgn_kthread = kthread_create(drgn_kthread_fn, (void *)0xF0F0F0F0, "drgn_kthread");
-	if (!drgn_kthread)
+	drgn_test_kthread = kthread_create(drgn_test_kthread_fn, NULL,
+					   "drgn_test_kthread");
+	if (!drgn_test_kthread)
 		return -1;
-	wake_up_process(drgn_kthread);
-	return 0;
+	wake_up_process(drgn_test_kthread);
+	wait_for_completion(&drgn_test_kthread_ready);
+	return kthread_park(drgn_test_kthread);
 }
 
 // Dummy function symbol.

tests/linux_kernel/test_stack_trace.py

@@ -17,17 +17,20 @@ from util import NORMALIZED_MACHINE_NAME
 
 
 class TestStackTrace(LinuxKernelTestCase):
-    def _assert_trace_in_sigwait(self, trace):
-        for frame in trace:
-            if frame.name and "sigtimedwait" in frame.name:
-                return
-        self.fail(f"sigwait frame not found in {str(trace)!r}")
+    def _test_drgn_test_kthread_trace(self, trace):
+        for i, frame in enumerate(trace):
+            if frame.name == "drgn_test_kthread_fn3":
+                break
+        else:
+            self.fail("Couldn't find drgn_test_kthread_fn3 frame")
+        self.assertEqual(trace[i + 1].name, "drgn_test_kthread_fn2")
+        self.assertEqual(trace[i + 2].name, "drgn_test_kthread_fn")
 
+    @skip_unless_have_test_kmod
     def test_by_task_struct(self):
-        with fork_and_sigwait() as pid:
-            self._assert_trace_in_sigwait(
-                self.prog.stack_trace(find_task(self.prog, pid))
-            )
+        self._test_drgn_test_kthread_trace(
+            self.prog.stack_trace(self.prog["drgn_test_kthread"])
+        )
 
     def _test_by_pid(self, orc):
         old_orc = int(os.environ.get("DRGN_PREFER_ORC_UNWINDER", "0")) != 0
@@ -38,9 +41,11 @@ class TestStackTrace(LinuxKernelTestCase):
                 prog = Program()
                 prog.set_kernel()
                 self._load_debug_info(prog)
-            with fork_and_sigwait() as pid:
-                self._assert_trace_in_sigwait(prog.stack_trace(pid))
+            self._test_drgn_test_kthread_trace(
+                self.prog.stack_trace(self.prog["drgn_test_kthread"].pid)
+            )
 
+    @skip_unless_have_test_kmod
     def test_by_pid_dwarf(self):
         self._test_by_pid(False)
 
@@ -48,22 +53,31 @@ class TestStackTrace(LinuxKernelTestCase):
         NORMALIZED_MACHINE_NAME == "x86_64",
         f"{NORMALIZED_MACHINE_NAME} does not use ORC",
     )
+    @skip_unless_have_test_kmod
     def test_by_pid_orc(self):
         self._test_by_pid(True)
 
+    @skip_unless_have_test_kmod
     def test_local_variable(self):
-        with fork_and_sigwait() as pid:
-            for frame in self.prog.stack_trace(pid):
-                if frame.name in ("context_switch", "__schedule"):
-                    try:
-                        prev = frame["prev"]
-                    except KeyError:
-                        continue
-                    if not prev.absent_:
-                        self.assertEqual(prev.pid, pid)
-                        break
-            else:
-                self.skipTest("prev not found in context_switch or __schedule")
+        for frame in self.prog.stack_trace(self.prog["drgn_test_kthread"]):
+            if frame.name == "drgn_test_kthread_fn3":
+                break
+        else:
+            self.fail("Couldn't find drgn_test_kthread_fn3 frame")
+        self.assertEqual(frame["a"], 1)
+        self.assertEqual(frame["b"], 2)
+        self.assertEqual(frame["c"], 3)
+
+    @skip_unless_have_test_kmod
+    def test_locals(self):
+        task = self.prog["drgn_test_kthread"]
+        stack_trace = self.prog.stack_trace(task)
+        for frame in stack_trace:
+            if frame.name == "drgn_test_kthread_fn3":
+                self.assertSetEqual(set(frame.locals()), {"a", "b", "c"})
+                break
+        else:
+            self.fail("Couldn't find drgn_test_kthread_fn3 frame")
 
     def test_pt_regs(self):
         # This won't unwind anything useful, but at least make sure it accepts
@@ -99,17 +113,3 @@ class TestStackTrace(LinuxKernelTestCase):
             assertReprPrettyEqualsStr(trace)
             for frame in trace:
                 assertReprPrettyEqualsStr(frame)
-
-    @skip_unless_have_test_kmod
-    def test_stack_locals(self):
-        task = self.prog["drgn_kthread"]
-        stack_trace = self.prog.stack_trace(task)
-        for frame in stack_trace:
-            if frame.symbol().name == "drgn_kthread_fn":
-                self.assertSetEqual(
-                    {"arg", "a", "b", "c"},
-                    set(frame.locals()),
-                )
-                break
-        else:
-            self.fail("Couldn't find drgn_kthread_fn frame")