helpers: workqueue: Fix errors seen for newer (v6.4 and later) kernels.

Signed-off-by: Imran Khan <[email protected]>

Author: imran-kn

drgn/helpers/linux/workqueue.py

@@ -11,15 +11,15 @@
 
 from typing import Iterator, Optional, Union
 
-from drgn import NULL, Object, Program, cast
+from drgn import NULL, IntegerLike, Object, Program, cast
 from drgn.helpers.common.format import escape_ascii_string
 from drgn.helpers.linux.idr import idr_find, idr_for_each
 from drgn.helpers.linux.list import (
     hlist_for_each_entry,
     list_empty,
     list_for_each_entry,
 )
-from drgn.helpers.linux.percpu import per_cpu
+from drgn.helpers.linux.percpu import per_cpu, per_cpu_ptr
 from drgn.helpers.linux.pid import find_task
 
 __all__ = (
@@ -43,19 +43,71 @@
     "show_one_worker_pool",
     "is_task_a_worker",
     "find_worker_executing_work",
+    "workqueue_get_pwq",
 )
 
 
 _PF_WQ_WORKER = 0x00000020
 
 
+def _work_offq_pool_none(prog: Program) -> IntegerLike:
+    # Linux kernel commit afa4bb778e48 ("workqueue: clean up WORK_* constant
+    # types, clarify masking") (in v6.4) changed WORK_OFFQ_POOL_NONE from
+    # constants of type enum to constants of type unsigned long.
+    try:
+        val = prog["WORK_OFFQ_POOL_NONE"].value_()
+    except KeyError:
+        val = (
+            Object(prog, "unsigned long", 1).value_() << prog["WORK_OFFQ_POOL_BITS"]
+        ) - 1
+    return val
+
+
+def _work_struct_wq_data_mask(prog: Program) -> IntegerLike:
+    # Linux kernel commit afa4bb778e48 ("workqueue: clean up WORK_* constant
+    # types, clarify masking") (in v6.4) changed WORK_STRUCT_WQ_DATA_MASK from
+    # constants of type enum to constants of type unsigned long.
+    try:
+        val = prog["WORK_STRUCT_WQ_DATA_MASK"].value_()
+    except KeyError:
+        val = ~((Object(prog, "unsigned long", 1) << prog["WORK_STRUCT_FLAG_BITS"]) - 1)
+    return val
+
+
 def _print_work(work: Object) -> None:
     prog = work.prog_
     print(
         f"        ({work.type_.type_name()})0x{work.value_():x}: func: {prog.symbol(work.func.value_()).name}"
     )
 
 
+def workqueue_get_pwq(workqueue: Object, cpu: int) -> Object:
+    """
+    Find pool_workqueue of a bound workqueue for a given CPU.
+
+    :param workqueue: ``struct workqueue_struct *``
+    :return: ``struct pool_workqueue *``.
+    """
+    # At first Linux kernel commit ee1ceef72754 ("workqueue: Rename wq->cpu_pwqs to
+    # wq->cpu_pwq") (in v6.6) renamed cpu_pwqs to cpu_pwq and then Linux kernel commit
+    # 687a9aa56f81("workqueue: Make per-cpu pool_workqueues allocated and released
+    # like unbound ones") (in v6.6) changed cpu_pwq to double pointer.
+    # As both of the changes were made in v6.6, there are no kernel versions
+    # with wq->cpu_pwq as a pointer. Still I have mentioned both the changes so that
+    # we can track both name change and type change of this member.
+    try:
+        pwq = per_cpu_ptr(workqueue.cpu_pwqs, cpu)
+    except AttributeError:
+        prog = workqueue.prog_
+        pwq = Object(
+            prog,
+            "struct pool_workqueue",
+            address=prog.read_word(per_cpu_ptr(workqueue.cpu_pwq, cpu)),
+        ).address_of_()
+
+    return pwq
+
+
 def for_each_workqueue(prog: Program) -> Iterator[Object]:
     """
     Iterate over all workqueues in the system.
@@ -210,9 +262,7 @@ def get_work_pwq(work: Object) -> Object:
     prog = work.prog_
     data = cast("unsigned long", work.data.counter.read_())
     if data & prog["WORK_STRUCT_PWQ"].value_():
-        return cast(
-            "struct pool_workqueue *", data & prog["WORK_STRUCT_WQ_DATA_MASK"].value_()
-        )
+        return cast("struct pool_workqueue *", data & _work_struct_wq_data_mask(prog))
     else:
         return NULL(work.prog_, "struct pool_workqueue *")
 
@@ -230,12 +280,12 @@ def get_work_pool(work: Object) -> Object:
     data = cast("unsigned long", work.data.counter.read_())
 
     if data & prog["WORK_STRUCT_PWQ"].value_():
-        pwq = data & prog["WORK_STRUCT_WQ_DATA_MASK"].value_()
+        pwq = data & _work_struct_wq_data_mask(prog)
         pool = Object(prog, "struct pool_workqueue", address=pwq).pool
     else:
         pool_id = data >> prog["WORK_OFFQ_POOL_SHIFT"].value_()
 
-        if pool_id == prog["WORK_OFFQ_POOL_NONE"].value_():
+        if pool_id == _work_offq_pool_none(prog):
             return NULL(work.prog_, "struct worker_pool *")
 
         pool = idr_find(prog["worker_pool_idr"].address_of_(), pool_id)

tests/linux_kernel/helpers/test_workqueue.py

@@ -19,6 +19,7 @@
     for_each_workqueue,
     get_work_pool,
     get_work_pwq,
+    workqueue_get_pwq,
 )
 from tests.linux_kernel import LinuxKernelTestCase, skip_unless_have_test_kmod
 
@@ -114,11 +115,12 @@ def test_for_each_pwq(self):
         # Since "drgn_test_wq" is a bound workqueue, list pwqs
         # should contain only per-cpu pwqs i.e cpu_pwqs
         pwqs = [pwq.value_() for pwq in for_each_pwq(self.test_wq)]
-        cpu_pwqs = [
-            per_cpu_ptr(self.test_wq.cpu_pwqs, cpu).value_()
+        cpu_pwqs_attr = "cpu_pwqs" if hasattr(self.test_wq, "cpu_pwqs") else "cpu_pwq"
+        cpu_pwqs_list = [
+            per_cpu_ptr(getattr(self.test_wq, cpu_pwqs_attr), cpu).value_()
             for cpu in for_each_online_cpu(self.prog)
         ]
-        self.assertEqual(pwqs.sort(), cpu_pwqs.sort())
+        self.assertEqual(pwqs.sort(), cpu_pwqs_list.sort())
 
     @skip_unless_have_test_kmod
     def test_for_each_pending_work(self):
@@ -142,16 +144,17 @@ def test_for_each_pending_work_in_pool(self):
 
     @skip_unless_have_test_kmod
     def test_for_each_pending_work_of_pwq(self):
-        cpu_pwqs_0 = per_cpu_ptr(self.test_wq.cpu_pwqs, 0)
+        cpu_pwqs_0 = workqueue_get_pwq(self.test_wq, 0)
         all_works_of_pwq = [
             work.value_() for work in for_each_pending_work_of_pwq(cpu_pwqs_0)
         ]
         self.assertEqual(all_works_of_pwq, self.test_works)
 
     @skip_unless_have_test_kmod
     def test_get_work_pwq(self):
-        cpu_pwqs_0 = per_cpu_ptr(self.test_wq.cpu_pwqs, 0)
-        cpu_pwqs_1 = per_cpu_ptr(self.test_wq.cpu_pwqs, 1)
+        cpu_pwqs_0 = workqueue_get_pwq(self.test_wq, 0)
+        cpu_pwqs_1 = workqueue_get_pwq(self.test_wq, 1)
+
         pwq = get_work_pwq(self.test_work_0)
         self.assertEqual(pwq, cpu_pwqs_0)
         self.assertNotEqual(pwq, cpu_pwqs_1)