Rewrite linux helper iterators in C

In preparation for introducing an API to represent threads, the linux
helper iterators, radix_tree_for_each, idr_for_each, for_each_pid, and
for_each_task have been rewritten in C. This will allow them to be
accessed from libdrgn, which will be necessary for the threads API.

Signed-off-by: Kevin Svetlitski <svetlitski@fb.com>

_drgn.pyi

@@ -2219,3 +2219,41 @@ def _linux_helper_kaslr_offset(prog: Program) -> int:
 def _linux_helper_pgtable_l5_enabled(prog: Program) -> bool:
     """Return whether 5-level paging is enabled."""
     ...
+
+def _linux_helper_radix_tree_for_each(root: Object) -> Iterator[Tuple[int, Object]]:
+    """
+    Iterate over all of the entries in a radix tree.
+
+    :param root: ``struct radix_tree_root *``
+    :return: Iterator of (index, ``void *``) tuples.
+    """
+    ...
+
+def _linux_helper_idr_for_each(idr: Object) -> Iterator[Tuple[int, Object]]:
+    """
+    Iterate over all of the entries in an IDR.
+
+    :param idr: ``struct idr *``
+    :return: Iterator of (index, ``void *``) tuples.
+    """
+    ...
+
+def _linux_helper_for_each_pid(prog_or_ns: Union[Program, Object]) -> Iterator[Object]:
+    """
+    Iterate over all PIDs in a namespace.
+
+    :param prog_or_ns: ``struct pid_namespace *`` to iterate over, or
+        :class:`Program` to iterate over initial PID namespace.
+    :return: Iterator of ``struct pid *`` objects.
+    """
+    ...
+
+def _linux_helper_for_each_task(prog_or_ns: Union[Program, Object]) -> Iterator[Object]:
+    """
+    Iterate over all of the tasks visible in a namespace.
+
+    :param prog_or_ns: ``struct pid_namespace *`` to iterate over, or
+        :class:`Program` to iterate over initial PID namespace.
+    :return: Iterator of ``struct task_struct *`` objects.
+    """
+    ...

drgn/helpers/linux/idr.py

@@ -11,28 +11,12 @@ an ID to a pointer. This currently only supports Linux v4.11+; before this,
 IDRs were not based on radix trees.
 """
 
-from typing import Iterator, Tuple
+from _drgn import (
-
+    _linux_helper_idr_find as idr_find,
-from _drgn import _linux_helper_idr_find as idr_find
+    _linux_helper_idr_for_each as idr_for_each,
-from drgn import Object
+)
-from drgn.helpers.linux.radixtree import radix_tree_for_each
 
 __all__ = (
     "idr_find",
     "idr_for_each",
 )
-
-
-def idr_for_each(idr: Object) -> Iterator[Tuple[int, Object]]:
-    """
-    Iterate over all of the entries in an IDR.
-
-    :param idr: ``struct idr *``
-    :return: Iterator of (index, ``void *``) tuples.
-    """
-    try:
-        base = idr.idr_base.value_()
-    except AttributeError:
-        base = 0
-    for index, entry in radix_tree_for_each(idr.idr_rt.address_of_()):
-        yield index + base, entry

drgn/helpers/linux/pid.py

@@ -9,16 +9,13 @@ The ``drgn.helpers.linux.pid`` module provides helpers for looking up process
 IDs and processes.
 """
 
-from typing import Iterator, Union
-
 from _drgn import (
     _linux_helper_find_pid as find_pid,
     _linux_helper_find_task as find_task,
+    _linux_helper_for_each_pid as for_each_pid,
+    _linux_helper_for_each_task as for_each_task,
     _linux_helper_pid_task as pid_task,
 )
-from drgn import Object, Program, cast, container_of
-from drgn.helpers.linux.idr import idr_for_each
-from drgn.helpers.linux.list import hlist_for_each_entry
 
 __all__ = (
     "find_pid",
@@ -27,49 +24,3 @@ __all__ = (
     "for_each_task",
     "pid_task",
 )
-
-
-def for_each_pid(prog_or_ns: Union[Program, Object]) -> Iterator[Object]:
-    """
-    Iterate over all PIDs in a namespace.
-
-    :param prog_or_ns: ``struct pid_namespace *`` to iterate over, or
-        :class:`Program` to iterate over initial PID namespace.
-    :return: Iterator of ``struct pid *`` objects.
-    """
-    if isinstance(prog_or_ns, Program):
-        prog = prog_or_ns
-        ns = prog_or_ns["init_pid_ns"].address_of_()
-    else:
-        prog = prog_or_ns.prog_
-        ns = prog_or_ns
-    if hasattr(ns, "idr"):
-        for nr, entry in idr_for_each(ns.idr):
-            yield cast("struct pid *", entry)
-    else:
-        pid_hash = prog["pid_hash"]
-        for i in range(1 << prog["pidhash_shift"].value_()):
-            for upid in hlist_for_each_entry(
-                "struct upid", pid_hash[i].address_of_(), "pid_chain"
-            ):
-                if upid.ns == ns:
-                    yield container_of(upid, "struct pid", f"numbers[{int(ns.level)}]")
-
-
-def for_each_task(prog_or_ns: Union[Program, Object]) -> Iterator[Object]:
-    """
-    Iterate over all of the tasks visible in a namespace.
-
-    :param prog_or_ns: ``struct pid_namespace *`` to iterate over, or
-        :class:`Program` to iterate over initial PID namespace.
-    :return: Iterator of ``struct task_struct *`` objects.
-    """
-    if isinstance(prog_or_ns, Program):
-        prog = prog_or_ns
-    else:
-        prog = prog_or_ns.prog_
-    PIDTYPE_PID = prog["PIDTYPE_PID"].value_()
-    for pid in for_each_pid(prog_or_ns):
-        task = pid_task(pid, PIDTYPE_PID)
-        if task:
-            yield task

drgn/helpers/linux/radixtree.py

@@ -9,54 +9,12 @@ The ``drgn.helpers.linux.radixtree`` module provides helpers for working with
 radix trees from :linux:`include/linux/radix-tree.h`.
 """
 
-from typing import Iterator, Tuple
+from _drgn import (
-
+    _linux_helper_radix_tree_for_each as radix_tree_for_each,
-from _drgn import _linux_helper_radix_tree_lookup as radix_tree_lookup
+    _linux_helper_radix_tree_lookup as radix_tree_lookup,
-from drgn import Object, cast
+)
 
 __all__ = (
     "radix_tree_for_each",
     "radix_tree_lookup",
 )
-
-_RADIX_TREE_ENTRY_MASK = 3
-
-
-def _is_internal_node(node: Object, internal_node: int) -> bool:
-    return (node.value_() & _RADIX_TREE_ENTRY_MASK) == internal_node
-
-
-def _entry_to_node(node: Object, internal_node: int) -> Object:
-    return Object(node.prog_, node.type_, value=node.value_() & ~internal_node)
-
-
-def _radix_tree_root_node(root: Object) -> Tuple[Object, int]:
-    try:
-        node = root.xa_head
-    except AttributeError:
-        return root.rnode.read_(), 1
-    else:
-        return cast("struct xa_node *", node).read_(), 2
-
-
-def radix_tree_for_each(root: Object) -> Iterator[Tuple[int, Object]]:
-    """
-    Iterate over all of the entries in a radix tree.
-
-    :param root: ``struct radix_tree_root *``
-    :return: Iterator of (index, ``void *``) tuples.
-    """
-    node, RADIX_TREE_INTERNAL_NODE = _radix_tree_root_node(root)
-
-    def aux(node: Object, index: int) -> Iterator[Tuple[int, Object]]:
-        if _is_internal_node(node, RADIX_TREE_INTERNAL_NODE):
-            parent = _entry_to_node(node, RADIX_TREE_INTERNAL_NODE)
-            for i, slot in enumerate(parent.slots):
-                yield from aux(
-                    cast(parent.type_, slot).read_(),
-                    index + (i << parent.shift.value_()),
-                )
-        elif node:
-            yield index, cast("void *", node)
-
-    yield from aux(node, 0)

libdrgn/helpers.h

@@ -15,6 +15,8 @@
 
 #include <stddef.h>
 #include <stdint.h>
+#include "drgn.h"
+#include "vector.h"
 
 struct drgn_object;
 struct drgn_program;
@@ -43,4 +45,100 @@ struct drgn_error *linux_helper_find_task(struct drgn_object *res,
 					  const struct drgn_object *ns,
 					  uint64_t pid);
 
+/**
+ * Iterator convention:
+ *
+ * For all of the iterators defined below, the convention for each of the
+ * `*_next` functions is that upon returning, `*ret` will point to space
+ * allocated inside of `iter`. The caller is free to do what they wish with
+ * this return value, but should note that it will be overwritten the next time
+ * the `*_next` function is called.
+ */
+
+DEFINE_VECTOR_TYPE(linux_helper_radix_tree_iter_frame_vector,
+		   struct linux_helper_radix_tree_iter_frame)
+
+struct linux_helper_radix_tree_iter_entry {
+	uint64_t index;
+	struct drgn_object node;
+};
+
+struct linux_helper_radix_tree_iter {
+	bool started;
+	struct drgn_object root;
+	// Current value to be yielded
+	struct linux_helper_radix_tree_iter_entry entry;
+	// We need this for later initialization of `drgn_object`s
+	struct drgn_program *prog;
+	// Frames to keep track of generator state
+	struct linux_helper_radix_tree_iter_frame_vector frames;
+	// One-time setup values that are persistent
+	uint64_t RADIX_TREE_INTERNAL_NODE;
+	uint64_t RADIX_TREE_MAP_MASK;
+	struct drgn_qualified_type node_type;
+};
+
+struct drgn_error *linux_helper_radix_tree_iter_init(struct linux_helper_radix_tree_iter *iter,
+						     const struct drgn_object *root);
+
+void linux_helper_radix_tree_iter_deinit(struct linux_helper_radix_tree_iter *iter);
+
+struct drgn_error *linux_helper_radix_tree_iter_next(struct linux_helper_radix_tree_iter *iter,
+						     struct linux_helper_radix_tree_iter_entry **ret);
+
+struct linux_helper_idr_iter {
+	struct linux_helper_radix_tree_iter iter;
+	uint64_t base;
+};
+
+struct drgn_error *linux_helper_idr_iter_init(struct linux_helper_idr_iter *iter,
+					      const struct drgn_object *idr);
+
+void linux_helper_idr_iter_deinit(struct linux_helper_idr_iter *iter);
+
+struct drgn_error *linux_helper_idr_iter_next(struct linux_helper_idr_iter *iter,
+					      struct linux_helper_radix_tree_iter_entry **ret);
+
+struct linux_helper_pid_iter {
+	bool has_idr;
+	struct drgn_qualified_type pid_type;
+	union {
+		// if has_idr
+		struct linux_helper_idr_iter iter;
+		// else
+		struct {
+			struct drgn_qualified_type upid_type;
+			struct drgn_object pid_hash;
+			struct drgn_object pos; // a `struct hlist_node*`
+			struct drgn_object ns;
+			struct drgn_object entry; // Current value of the iterator
+			size_t index; // Current loop index
+			char member_specifier[sizeof("numbers[]") + 20];
+			// 20 = maximum length of a uint64_t as a string
+			// Space for the null terminator is included as part of the sizeof on the string literal
+		};
+	};
+};
+
+struct drgn_error *linux_helper_pid_iter_init(struct linux_helper_pid_iter *iter,
+					      const struct drgn_object *ns);
+
+void linux_helper_pid_iter_deinit(struct linux_helper_pid_iter *iter);
+
+struct drgn_error *linux_helper_pid_iter_next(struct linux_helper_pid_iter *iter,
+					      struct drgn_object **ret);
+
+struct linux_helper_task_iter {
+	struct linux_helper_pid_iter iter;
+	uint64_t PIDTYPE_PID;
+};
+
+struct drgn_error *linux_helper_task_iter_init(struct linux_helper_task_iter *iter,
+					       const struct drgn_object *ns);
+
+void linux_helper_task_iter_deinit(struct linux_helper_task_iter *iter);
+
+struct drgn_error *linux_helper_task_iter_next(struct linux_helper_task_iter *iter,
+					       struct drgn_object **ret);
+
 #endif /* DRGN_HELPERS_H */

libdrgn/linux_kernel_helpers.c

@@ -6,10 +6,13 @@
 #include <inttypes.h>
 
 #include "drgn.h"
+#include "helpers.h"
 #include "minmax.h"
 #include "platform.h"
 #include "program.h"
 
+static const uint64_t RADIX_TREE_ENTRY_MASK = 3;
+
 struct drgn_error *linux_helper_read_vm(struct drgn_program *prog,
 					uint64_t pgtable, uint64_t virt_addr,
 					void *buf, size_t count)
@@ -99,12 +102,67 @@ struct drgn_error *linux_helper_read_vm(struct drgn_program *prog,
 	return err;
 }
 
+static struct drgn_error *
+radix_tree_init(struct drgn_program *prog, const struct drgn_object *root,
+		uint64_t *RADIX_TREE_INTERNAL_NODE_ret,
+		uint64_t *RADIX_TREE_MAP_MASK_ret,
+		struct drgn_qualified_type *node_type_ret,
+		struct drgn_object *node_ret)
+{
+	struct drgn_error *err =
+		drgn_object_member_dereference(node_ret, root, "xa_head");
+	/* node = root->xa_head */
+	if (!err) {
+		err = drgn_program_find_type(prog, "struct xa_node *", NULL,
+					     node_type_ret);
+		if (err)
+			return err;
+		*RADIX_TREE_INTERNAL_NODE_ret = 2;
+	} else if (err->code == DRGN_ERROR_LOOKUP) {
+		drgn_error_destroy(err);
+		/* node = (void *)root.rnode */
+		err = drgn_object_member_dereference(node_ret, root, "rnode");
+		if (err)
+			return err;
+		err = drgn_program_find_type(prog, "void *", NULL,
+					     node_type_ret);
+		if (err)
+			return err;
+		err = drgn_object_cast(node_ret, *node_type_ret, node_ret);
+		if (err)
+			return err;
+		err = drgn_program_find_type(prog, "struct radix_tree_node *",
+					     NULL, node_type_ret);
+		if (err)
+			return err;
+		*RADIX_TREE_INTERNAL_NODE_ret = 1;
+	} else {
+		return err;
+	}
+
+	struct drgn_type_member *member;
+	uint64_t member_bit_offset;
+	err = drgn_type_find_member(drgn_type_type(node_type_ret->type).type,
+				    "slots", &member, &member_bit_offset);
+	if (err)
+		return err;
+	struct drgn_qualified_type member_type;
+	err = drgn_member_type(member, &member_type, NULL);
+	if (err)
+		return err;
+	if (drgn_type_kind(member_type.type) != DRGN_TYPE_ARRAY)
+		return drgn_error_create(
+			DRGN_ERROR_TYPE,
+			"struct radix_tree_node slots member is not an array");
+	*RADIX_TREE_MAP_MASK_ret = drgn_type_length(member_type.type) - 1;
+	return NULL;
+}
+
 struct drgn_error *
 linux_helper_radix_tree_lookup(struct drgn_object *res,
 			       const struct drgn_object *root, uint64_t index)
 {
 	struct drgn_error *err;
-	static const uint64_t RADIX_TREE_ENTRY_MASK = 3;
 	uint64_t RADIX_TREE_INTERNAL_NODE;
 	uint64_t RADIX_TREE_MAP_MASK;
 	struct drgn_object node, tmp;
@@ -112,55 +170,11 @@ linux_helper_radix_tree_lookup(struct drgn_object *res,
 
 	drgn_object_init(&node, drgn_object_program(res));
 	drgn_object_init(&tmp, drgn_object_program(res));
-
+	err = radix_tree_init(drgn_object_program(root), root,
-	/* node = root->xa_head */
+				  &RADIX_TREE_INTERNAL_NODE, &RADIX_TREE_MAP_MASK,
-	err = drgn_object_member_dereference(&node, root, "xa_head");
+				  &node_type, &node);
-	if (!err) {
-		err = drgn_program_find_type(drgn_object_program(res),
-					     "struct xa_node *", NULL,
-					     &node_type);
 	if (err)
 		goto out;
-		RADIX_TREE_INTERNAL_NODE = 2;
-	} else if (err->code == DRGN_ERROR_LOOKUP) {
-		drgn_error_destroy(err);
-		/* node = (void *)root.rnode */
-		err = drgn_object_member_dereference(&node, root, "rnode");
-		if (err)
-			goto out;
-		err = drgn_program_find_type(drgn_object_program(res), "void *",
-					     NULL, &node_type);
-		if (err)
-			goto out;
-		err = drgn_object_cast(&node, node_type, &node);
-		if (err)
-			goto out;
-		err = drgn_program_find_type(drgn_object_program(res),
-					     "struct radix_tree_node *", NULL,
-					     &node_type);
-		if (err)
-			goto out;
-		RADIX_TREE_INTERNAL_NODE = 1;
-	} else {
-		goto out;
-	}
-
-	struct drgn_type_member *member;
-	uint64_t member_bit_offset;
-	err = drgn_type_find_member(drgn_type_type(node_type.type).type,
-				    "slots", &member, &member_bit_offset);
-	if (err)
-		goto out;
-	struct drgn_qualified_type member_type;
-	err = drgn_member_type(member, &member_type, NULL);
-	if (err)
-		goto out;
-	if (drgn_type_kind(member_type.type) != DRGN_TYPE_ARRAY) {
-		err = drgn_error_create(DRGN_ERROR_TYPE,
-					"struct radix_tree_node slots member is not an array");
-		goto out;
-	}
-	RADIX_TREE_MAP_MASK = drgn_type_length(member_type.type) - 1;
 
 	for (;;) {
 		uint64_t value;
@@ -243,6 +257,36 @@ out:
 	return err;
 }
 
+static struct drgn_error *pid_hash_init(struct drgn_program *prog,
+					       const struct drgn_object *ns,
+					       struct drgn_qualified_type *upid_type_ret,
+					       uint64_t *pidhash_length_ret, uint64_t *ns_level_ret)
+{
+	struct drgn_error *err;
+	struct drgn_object ns_level, pidhash_shift;
+	drgn_object_init(&ns_level, prog);
+	drgn_object_init(&pidhash_shift, prog);
+	err = drgn_program_find_type(prog, "struct upid", NULL, upid_type_ret);
+	if (err)
+		goto out;
+	err = drgn_program_find_object(prog, "pidhash_shift", NULL, DRGN_FIND_OBJECT_ANY,
+				       &pidhash_shift);
+	if (err)
+		goto out;
+	err = drgn_object_read_unsigned(&pidhash_shift, pidhash_length_ret);
+	if (err)
+		goto out;
+	// *pidhash_length_ret = 1 << pidhash_shift
+	*pidhash_length_ret = *pidhash_length_ret >= 64 ? 0 : UINT64_C(1) << *pidhash_length_ret;
+	err = drgn_object_member_dereference(&ns_level, ns, "level");
+	if (err)
+		goto out;
+	err = drgn_object_read_unsigned(&ns_level, ns_level_ret);
+out:
+	drgn_object_deinit(&ns_level);
+	return err;
+}
+
 /*
  * Before Linux kernel commit 95846ecf9dac ("pid: replace pid bitmap
  * implementation with IDR API") (in v4.15), (struct pid_namespace).idr does not
@@ -257,15 +301,27 @@ find_pid_in_pid_hash(struct drgn_object *res, const struct drgn_object *ns,
 {
 	struct drgn_error *err;
 
+	struct drgn_object node, tmp;
+	drgn_object_init(&node, drgn_object_program(res));
+	drgn_object_init(&tmp, drgn_object_program(res));
+
+	err = drgn_object_read(&tmp, ns);
+	if (err)
+		goto out;
+	struct drgn_qualified_type upid_type;
+	uint64_t i, ns_level;
+	err = pid_hash_init(drgn_object_program(res), &tmp, &upid_type, &i,
+				   &ns_level);
+	if (err)
+		goto out;
 	struct drgn_qualified_type pidp_type;
-	err = drgn_program_find_type(drgn_object_program(res), "struct pid *",
+	err = drgn_program_find_type(drgn_object_program(res), "struct pid *", NULL,
-				     NULL, &pidp_type);
+					 &pidp_type);
 	if (err)
 		return err;
 
-	struct drgn_qualified_type upid_type;
+	uint64_t ns_addr;
-	err = drgn_program_find_type(drgn_object_program(res), "struct upid",
+	err = drgn_object_read_unsigned(&tmp, &ns_addr);
-				     NULL, &upid_type);
 	if (err)
 		return err;
 
@@ -298,40 +354,6 @@ find_pid_in_pid_hash(struct drgn_object *res, const struct drgn_object *ns,
 	if (err)
 		return err;
 
-	struct drgn_object node, tmp;
-	drgn_object_init(&node, drgn_object_program(res));
-	drgn_object_init(&tmp, drgn_object_program(res));
-
-	err = drgn_object_read(&tmp, ns);
-	if (err)
-		goto out;
-	uint64_t ns_addr;
-	err = drgn_object_read_unsigned(&tmp, &ns_addr);
-	if (err)
-		goto out;
-	union drgn_value ns_level;
-	err = drgn_object_member_dereference(&tmp, &tmp, "level");
-	if (err)
-		goto out;
-	err = drgn_object_read_integer(&tmp, &ns_level);
-	if (err)
-		goto out;
-
-	/* i = 1 << pidhash_shift */
-	err = drgn_program_find_object(drgn_object_program(res),
-				       "pidhash_shift", NULL,
-				       DRGN_FIND_OBJECT_ANY, &tmp);
-	if (err)
-		goto out;
-	union drgn_value pidhash_shift;
-	err = drgn_object_read_integer(&tmp, &pidhash_shift);
-	if (err)
-		goto out;
-	uint64_t i;
-	if (pidhash_shift.uvalue >= 64)
-		i = 0;
-	else
-		i = UINT64_C(1) << pidhash_shift.uvalue;
 	while (i--) {
 		/* for (node = pid_hash[i].first; node; node = node->next) */
 		err = drgn_object_subscript(&node, pid_hash, i);
@@ -382,7 +404,7 @@ find_pid_in_pid_hash(struct drgn_object *res, const struct drgn_object *ns,
 				goto next;
 
 			sprintf(member, "numbers[%" PRIu64 "].pid_chain",
-				ns_level.uvalue);
+				ns_level);
 			err = drgn_object_container_of(res, &node,
 						       drgn_type_type(pidp_type.type),
 						       member);
@@ -533,3 +555,381 @@ out:
 	drgn_object_deinit(&pid_obj);
 	return err;
 }
+
+struct linux_helper_radix_tree_iter_frame {
+	struct drgn_object slots;
+	uint64_t index;
+	uint64_t shift;
+	uint64_t next_slot;
+};
+
+DEFINE_VECTOR_FUNCTIONS(linux_helper_radix_tree_iter_frame_vector)
+
+struct drgn_error *linux_helper_radix_tree_iter_init(struct linux_helper_radix_tree_iter *iter,
+						     const struct drgn_object *root)
+{
+	struct drgn_program *prog = drgn_object_program(root);
+	iter->started = false;
+	drgn_object_init(&iter->root, prog);
+	drgn_object_init(&iter->entry.node, prog);
+	iter->entry.index = 0;
+	iter->prog = prog;
+
+	struct drgn_error *err =
+		radix_tree_init(prog, root, &iter->RADIX_TREE_INTERNAL_NODE,
+				&iter->RADIX_TREE_MAP_MASK, &iter->node_type, &iter->root);
+
+	if (err) {
+		drgn_object_deinit(&iter->root);
+		drgn_object_deinit(&iter->entry.node);
+		return err;
+	}
+
+	linux_helper_radix_tree_iter_frame_vector_init(&iter->frames);
+	return NULL;
+}
+
+void linux_helper_radix_tree_iter_deinit(struct linux_helper_radix_tree_iter *iter)
+{
+	drgn_object_deinit(&iter->root);
+	drgn_object_deinit(&iter->entry.node);
+	while (iter->frames.size) {
+		drgn_object_deinit(
+			&linux_helper_radix_tree_iter_frame_vector_pop(&iter->frames)->slots);
+	}
+	linux_helper_radix_tree_iter_frame_vector_deinit(&iter->frames);
+}
+
+static struct drgn_error *radix_tree_iter_handle_node(struct linux_helper_radix_tree_iter *iter,
+						      struct drgn_object *_node, uint64_t index,
+						      bool *entry_populated_ret)
+{
+	struct drgn_object *node = &iter->entry.node;
+	struct drgn_error *err;
+	uint64_t value;
+
+	err = drgn_object_read(node, _node);
+	if (err)
+		return err;
+	err = drgn_object_read_unsigned(node, &value);
+	if (err)
+		return err;
+	if ((value & RADIX_TREE_ENTRY_MASK) != iter->RADIX_TREE_INTERNAL_NODE) {
+		// Base-case, node is NOT internal
+		if (value) {
+			*entry_populated_ret = true;
+			iter->entry.index = index;
+		}
+		return NULL;
+	}
+
+	*entry_populated_ret = false;
+
+	// We are dealing with an internal node, and must iterate over its slots
+
+	err = drgn_object_set_unsigned(node, iter->node_type,
+				       value & ~iter->RADIX_TREE_INTERNAL_NODE, 0);
+	if (err)
+		return err;
+	struct linux_helper_radix_tree_iter_frame *frame =
+		linux_helper_radix_tree_iter_frame_vector_append_entry(&iter->frames);
+	if (!frame)
+		return &drgn_enomem;
+	frame->index = index;
+	frame->next_slot = 0;
+	drgn_object_init(&frame->slots, iter->prog);
+	// We temporarily use `frame->slots` to hold `shift` in order to avoid
+	// using another `struct drgn_object`.
+	err = drgn_object_member_dereference(&frame->slots, node, "shift");
+	if (err)
+		goto err_frame;
+	err = drgn_object_read_unsigned(&frame->slots, &frame->shift);
+	if (err)
+		goto err_frame;
+	// Now `frame->slots` is actually used for `slots`.
+	err = drgn_object_member_dereference(&frame->slots, node, "slots");
+	if (err)
+		goto err_frame;
+	return NULL;
+
+err_frame:
+	drgn_object_deinit(&frame->slots);
+	linux_helper_radix_tree_iter_frame_vector_pop(&iter->frames);
+	return err;
+}
+
+struct drgn_error *linux_helper_radix_tree_iter_next(struct linux_helper_radix_tree_iter *iter,
+						     struct linux_helper_radix_tree_iter_entry **ret)
+{
+	bool entry_populated = false;
+	struct drgn_error *err = NULL;
+	struct drgn_object node;
+	drgn_object_init(&node, iter->prog);
+	if (!iter->started) {
+		iter->started = true;
+		err = radix_tree_iter_handle_node(iter, &iter->root, 0, &entry_populated);
+	}
+
+	while (!err && !entry_populated && iter->frames.size) {
+		struct linux_helper_radix_tree_iter_frame *frame =
+			&iter->frames.data[iter->frames.size - 1];
+		if (frame->next_slot <= iter->RADIX_TREE_MAP_MASK) {
+			err = drgn_object_subscript(&node, &frame->slots, frame->next_slot);
+			if (!err)
+				err = radix_tree_iter_handle_node(iter, &node,
+								  frame->index + (frame->next_slot++
+										  << frame->shift),
+								  &entry_populated);
+		} else {
+			drgn_object_deinit(&frame->slots);
+			linux_helper_radix_tree_iter_frame_vector_pop(&iter->frames);
+		}
+	}
+	if (!err)
+		*ret = entry_populated ? &iter->entry : NULL;
+	drgn_object_deinit(&node);
+	return err;
+}
+
+struct drgn_error *linux_helper_idr_iter_init(struct linux_helper_idr_iter *iter,
+					      const struct drgn_object *idr)
+{
+	struct drgn_error *err;
+	struct drgn_object idr_rt, idr_base;
+	drgn_object_init(&idr_rt, drgn_object_program(idr));
+	drgn_object_init(&idr_base, drgn_object_program(idr));
+
+	err = drgn_object_member(&idr_base, idr, "idr_base");
+	if (!err) {
+		err = drgn_object_read_unsigned(&idr_base, &iter->base);
+		if (err)
+			goto out;
+	} else if (err->code == DRGN_ERROR_LOOKUP) {
+		drgn_error_destroy(err);
+		iter->base = 0;
+	} else {
+		goto out;
+	}
+
+	err = drgn_object_member(&idr_rt, idr, "idr_rt");
+	if (err)
+		goto out;
+	err = drgn_object_address_of(&idr_rt, &idr_rt);
+	if (err)
+		goto out;
+	err = linux_helper_radix_tree_iter_init(&iter->iter, &idr_rt);
+out:
+	drgn_object_deinit(&idr_rt);
+	drgn_object_deinit(&idr_base);
+	return err;
+}
+
+void linux_helper_idr_iter_deinit(struct linux_helper_idr_iter *iter)
+{
+	linux_helper_radix_tree_iter_deinit(&iter->iter);
+}
+
+struct drgn_error *linux_helper_idr_iter_next(struct linux_helper_idr_iter *iter,
+					      struct linux_helper_radix_tree_iter_entry **ret)
+{
+	struct drgn_error *err = linux_helper_radix_tree_iter_next(&iter->iter, ret);
+	if (!err && *ret)
+		(*ret)->index += iter->base;
+	return err;
+}
+
+// See `find_pid_in_pid_hash`
+static struct drgn_error *pid_iter_init_pid_hash(struct drgn_program *prog,
+						 const struct drgn_object *ns,
+						 struct linux_helper_pid_iter *iter)
+{
+	struct drgn_error *err;
+	drgn_object_init(&iter->pid_hash, prog);
+	drgn_object_init(&iter->pos, prog);
+	drgn_object_init(&iter->ns, prog);
+	drgn_object_init(&iter->entry, prog);
+	err = drgn_program_find_object(prog, "pid_hash", NULL, DRGN_FIND_OBJECT_VARIABLE,
+				       &iter->pid_hash);
+	if (err)
+		goto out;
+	struct drgn_qualified_type void_star_type;
+	err = drgn_program_find_type(prog, "void *", NULL, &void_star_type);
+	if (err)
+		goto out;
+	err = drgn_object_set_unsigned(&iter->pos, void_star_type, 0, 0);
+	if (err)
+		goto out;
+	err = drgn_object_copy(&iter->ns, ns);
+	if (err)
+		goto out;
+	uint64_t ns_level;
+	err = pid_hash_init(prog, ns, &iter->upid_type, &iter->index, &ns_level);
+	if (err)
+		goto out;
+	snprintf(iter->member_specifier, sizeof(iter->member_specifier), "numbers[%" PRIu64 "]",
+		 ns_level);
+	err = drgn_program_find_type(prog, "struct pid", NULL, &iter->pid_type);
+	if (err)
+		goto out;
+	err = drgn_program_find_type(prog, "struct upid", NULL, &iter->upid_type);
+	if (err)
+		goto out;
+out:
+	if (err) {
+		drgn_object_deinit(&iter->pid_hash);
+		drgn_object_deinit(&iter->pos);
+		drgn_object_deinit(&iter->ns);
+		drgn_object_deinit(&iter->entry);
+	}
+	return err;
+}
+
+struct drgn_error *linux_helper_pid_iter_init(struct linux_helper_pid_iter *iter,
+					      const struct drgn_object *ns)
+{
+	struct drgn_program *prog = drgn_object_program(ns);
+	struct drgn_error *err;
+	struct drgn_object idr;
+	drgn_object_init(&idr, prog);
+
+	err = drgn_object_member_dereference(&idr, ns, "idr");
+	if (!err) {
+		iter->has_idr = true;
+		err = drgn_program_find_type(prog, "struct pid *", NULL, &iter->pid_type);
+		if (!err)
+			err = linux_helper_idr_iter_init(&iter->iter, &idr);
+	} else if (err->code == DRGN_ERROR_LOOKUP) {
+		iter->has_idr = false;
+		drgn_error_destroy(err);
+		err = pid_iter_init_pid_hash(prog, ns, iter);
+	}
+
+	drgn_object_deinit(&idr);
+	return err;
+}
+
+void linux_helper_pid_iter_deinit(struct linux_helper_pid_iter *iter)
+{
+	if (iter->has_idr) {
+		linux_helper_idr_iter_deinit(&iter->iter);
+	} else {
+		drgn_object_deinit(&iter->pid_hash);
+		drgn_object_deinit(&iter->pos);
+		drgn_object_deinit(&iter->ns);
+		drgn_object_deinit(&iter->entry);
+	}
+}
+
+struct drgn_error *linux_helper_pid_iter_next(struct linux_helper_pid_iter *iter,
+					      struct drgn_object **ret)
+{
+	if (iter->has_idr) {
+		struct linux_helper_radix_tree_iter_entry *entry;
+		struct drgn_error *err = linux_helper_idr_iter_next(&iter->iter, &entry);
+		if (err)
+			return err;
+		if (!entry) {
+			*ret = NULL;
+			return NULL;
+		}
+		err = drgn_object_cast(&entry->node, iter->pid_type, &entry->node);
+		if (!err)
+			*ret = &entry->node;
+		return err;
+	}
+
+	struct drgn_error *err = NULL;
+	struct drgn_object upid, upid_ns;
+	drgn_object_init(&upid, drgn_object_program(&iter->ns));
+	drgn_object_init(&upid_ns, drgn_object_program(&iter->ns));
+
+	for (;;) {
+		for (;;) {
+			bool is_truthy;
+			err = drgn_object_bool(&iter->pos, &is_truthy);
+			if (err)
+				goto out;
+			if (is_truthy)
+				break;
+			if (iter->index == 0) {
+				*ret = NULL;
+				goto out;
+			}
+			err = drgn_object_subscript(&iter->pos, &iter->pid_hash, --iter->index);
+			if (err)
+				goto out;
+			err = drgn_object_member(&iter->pos, &iter->pos, "first");
+			if (err)
+				goto out;
+			err = drgn_object_bool(&iter->pos, &is_truthy);
+			if (err)
+				goto out;
+		}
+		err = drgn_object_container_of(&upid, &iter->pos, iter->upid_type, "pid_chain");
+		if (err)
+			goto out;
+		err = drgn_object_member_dereference(&iter->pos, &iter->pos, "next");
+		if (err)
+			goto out;
+		err = drgn_object_member_dereference(&upid_ns, &upid, "ns");
+		if (err)
+			goto out;
+		int ns_cmp_result;
+		err = drgn_object_cmp(&upid_ns, &iter->ns, &ns_cmp_result);
+		if (err)
+			goto out;
+		if (ns_cmp_result == 0) {
+			err = drgn_object_container_of(&iter->entry, &upid, iter->pid_type,
+						       iter->member_specifier);
+			if (!err)
+				*ret = &iter->entry;
+			goto out;
+		}
+	}
+
+out:
+	drgn_object_deinit(&upid);
+	drgn_object_deinit(&upid_ns);
+	return err;
+}
+
+struct drgn_error *linux_helper_task_iter_init(struct linux_helper_task_iter *iter,
+					       const struct drgn_object *ns)
+{
+	struct drgn_program *prog = drgn_object_program(ns);
+	struct drgn_error *err = linux_helper_pid_iter_init(&iter->iter, ns);
+	if (err)
+		return err;
+	struct drgn_object PIDTYPE_PID;
+	drgn_object_init(&PIDTYPE_PID, prog);
+	err = drgn_program_find_object(prog, "PIDTYPE_PID", NULL, DRGN_FIND_OBJECT_CONSTANT,
+				       &PIDTYPE_PID);
+	if (!err)
+		err = drgn_object_read_unsigned(&PIDTYPE_PID, &iter->PIDTYPE_PID);
+	if (err)
+		linux_helper_pid_iter_deinit(&iter->iter);
+	drgn_object_deinit(&PIDTYPE_PID);
+	return err;
+}
+
+struct drgn_error *linux_helper_task_iter_next(struct linux_helper_task_iter *iter,
+					       struct drgn_object **ret)
+{
+	struct drgn_error *err;
+	bool value_is_truthy;
+	do {
+		err = linux_helper_pid_iter_next(&iter->iter, ret);
+		if (err || !*ret)
+			return err;
+		err = linux_helper_pid_task(*ret, *ret, iter->PIDTYPE_PID);
+		if (err)
+			return err;
+		err = drgn_object_bool(*ret, &value_is_truthy);
+	} while (!err && !value_is_truthy);
+	return err;
+}
+
+void linux_helper_task_iter_deinit(struct linux_helper_task_iter *iter)
+{
+	linux_helper_pid_iter_deinit(&iter->iter);
+}

libdrgn/python/drgnpy.h

@@ -97,6 +97,14 @@ typedef struct {
 	struct pyobjectp_set objects;
 } Program;
 
+typedef struct _GenericIterator {
+	PyObject_HEAD
+	Program *prog;
+	void *iter;
+	PyObject *(*next)(struct _GenericIterator *);
+	void (*iter_deinit)(void *);
+} GenericIterator;
+
 typedef struct {
 	PyObject_HEAD
 	const struct drgn_register *reg;
@@ -167,6 +175,7 @@ extern PyObject *TypeKind_class;
 extern PyTypeObject DrgnObject_type;
 extern PyTypeObject DrgnType_type;
 extern PyTypeObject FaultError_type;
+extern PyTypeObject GenericIterator_type;
 extern PyTypeObject Language_type;
 extern PyTypeObject ObjectIterator_type;
 extern PyTypeObject Platform_type;
@@ -297,5 +306,17 @@ PyObject *drgnpy_linux_helper_kaslr_offset(PyObject *self, PyObject *args,
 					   PyObject *kwds);
 PyObject *drgnpy_linux_helper_pgtable_l5_enabled(PyObject *self, PyObject *args,
 						 PyObject *kwds);
+GenericIterator *drgnpy_linux_helper_for_each_task(PyObject *self,
+						      PyObject *args,
+						      PyObject *kwds);
+GenericIterator *drgnpy_linux_helper_for_each_pid(PyObject *self,
+						     PyObject *args,
+						     PyObject *kwds);
+GenericIterator *drgnpy_linux_helper_idr_for_each(PyObject *self,
+						     PyObject *args,
+						     PyObject *kwds);
+GenericIterator *drgnpy_linux_helper_radix_tree_for_each(PyObject *self,
+							    PyObject *args,
+							    PyObject *kwds);
 
 #endif /* DRGNPY_H */

libdrgn/python/helpers.c

@@ -249,3 +249,251 @@ PyObject *drgnpy_linux_helper_pgtable_l5_enabled(PyObject *self, PyObject *args,
 		return PyErr_Format(PyExc_ValueError, "not Linux kernel");
 	Py_RETURN_BOOL(prog->prog.vmcoreinfo.pgtable_l5_enabled);
 }
+
+static void GenericIterator_dealloc(GenericIterator *self)
+{
+	if (self->iter) {
+		self->iter_deinit(self->iter);
+		free(self->iter);
+	}
+	Py_XDECREF(self->prog);
+	Py_TYPE(self)->tp_free((PyObject *)self);
+}
+
+static PyObject *GenericIterator_next(GenericIterator *self)
+{
+	return self->next(self);
+}
+
+PyTypeObject GenericIterator_type = {
+	PyVarObject_HEAD_INIT(NULL, 0)
+	.tp_name = "_drgn._GenericIterator",
+	.tp_basicsize = sizeof(GenericIterator),
+	.tp_dealloc = (destructor)GenericIterator_dealloc,
+	.tp_flags = Py_TPFLAGS_DEFAULT,
+	.tp_iter = PyObject_SelfIter,
+	.tp_iternext = (iternextfunc)GenericIterator_next,
+};
+
+static PyObject *for_each_task_next(GenericIterator *self)
+{
+	struct drgn_error *err;
+	struct drgn_object *entry;
+	err = linux_helper_task_iter_next(self->iter, &entry);
+	if (err)
+		return set_drgn_error(err);
+	if (!entry)
+		return NULL;
+	DrgnObject *ret = DrgnObject_alloc(self->prog);
+	if (!ret)
+		return NULL;
+	err = drgn_object_copy(&ret->obj, entry);
+	if (err) {
+		Py_DECREF(ret);
+		return set_drgn_error(err);
+	}
+	return (PyObject *)ret;
+}
+
+GenericIterator *drgnpy_linux_helper_for_each_task(PyObject *self,
+						   PyObject *args,
+						   PyObject *kwds)
+{
+	static char *keywords[] = {"prog_or_ns", NULL};
+	struct drgn_error *err = NULL;
+	struct prog_or_ns_arg prog_or_ns;
+
+	if (!PyArg_ParseTupleAndKeywords(args, kwds, "O&:for_each_task",
+					 keywords, &prog_or_pid_ns_converter,
+					 &prog_or_ns))
+		return NULL;
+
+	GenericIterator *iterator =
+		(GenericIterator *)GenericIterator_type.tp_alloc(
+			&GenericIterator_type, 0);
+	if (!iterator)
+		goto out;
+	iterator->prog = prog_or_ns.prog;
+	Py_INCREF(iterator->prog);
+	iterator->next = for_each_task_next;
+	iterator->iter_deinit = (void (*)(void *))linux_helper_task_iter_deinit;
+	iterator->iter = malloc(sizeof(struct linux_helper_task_iter));
+	if (!iterator->iter) {
+		PyErr_NoMemory();
+		Py_DECREF(iterator);
+		iterator = NULL;
+		goto out;
+	}
+	err = linux_helper_task_iter_init(iterator->iter, prog_or_ns.ns);
+	if (err) {
+		set_drgn_error(err);
+		Py_DECREF(iterator);
+		iterator = NULL;
+	}
+out:
+	prog_or_ns_cleanup(&prog_or_ns);
+	return iterator;
+}
+
+static PyObject *for_each_pid_next(GenericIterator *self)
+{
+	struct drgn_error *err;
+	struct drgn_object *entry;
+	err = linux_helper_pid_iter_next(self->iter, &entry);
+	if (err)
+		return set_drgn_error(err);
+	if (!entry)
+		return NULL;
+	DrgnObject *ret = DrgnObject_alloc(self->prog);
+	if (!ret)
+		return NULL;
+	err = drgn_object_copy(&ret->obj, entry);
+	if (err) {
+		Py_DECREF(ret);
+		return set_drgn_error(err);
+	}
+	return (PyObject *)ret;
+}
+
+GenericIterator *
+drgnpy_linux_helper_for_each_pid(PyObject *self, PyObject *args, PyObject *kwds)
+{
+	static char *keywords[] = {"prog_or_ns", NULL};
+	struct drgn_error *err = NULL;
+	struct prog_or_ns_arg prog_or_ns;
+
+	if (!PyArg_ParseTupleAndKeywords(args, kwds, "O&:for_each_pid",
+					 keywords, &prog_or_pid_ns_converter,
+					 &prog_or_ns))
+		return NULL;
+
+	GenericIterator *iterator =
+		(GenericIterator *)GenericIterator_type.tp_alloc(
+			&GenericIterator_type, 0);
+	if (!iterator)
+		goto out;
+	iterator->prog = prog_or_ns.prog;
+	Py_INCREF(iterator->prog);
+	iterator->next = for_each_pid_next;
+	iterator->iter_deinit = (void (*)(void *))linux_helper_pid_iter_deinit;
+	iterator->iter = malloc(sizeof(struct linux_helper_pid_iter));
+	if (!iterator->iter) {
+		PyErr_NoMemory();
+		Py_DECREF(iterator);
+		iterator = NULL;
+		goto out;
+	}
+	err = linux_helper_pid_iter_init(iterator->iter, prog_or_ns.ns);
+	if (err) {
+		set_drgn_error(err);
+		Py_DECREF(iterator);
+		iterator = NULL;
+	}
+out:
+	prog_or_ns_cleanup(&prog_or_ns);
+	return iterator;
+}
+
+static PyObject *idr_iter_entry_wrap(struct linux_helper_radix_tree_iter_entry *entry,
+				     Program *prog)
+{
+	DrgnObject *node = DrgnObject_alloc(prog);
+	if (!node)
+		return NULL;
+	struct drgn_error *err = drgn_object_copy(&node->obj, &entry->node);
+	if (err) {
+		Py_DECREF(node);
+		return set_drgn_error(err);
+	}
+	PyObject *ret =
+		Py_BuildValue("KO", (unsigned long long)entry->index, node);
+	Py_DECREF(node);
+	return ret;
+}
+
+static PyObject *idr_for_each_next(GenericIterator *self)
+{
+	struct linux_helper_radix_tree_iter_entry *entry;
+	struct drgn_error *err = linux_helper_idr_iter_next(self->iter, &entry);
+	if (err)
+		return set_drgn_error(err);
+	if (!entry)
+		return NULL;
+	return idr_iter_entry_wrap(entry, self->prog);
+}
+
+GenericIterator *
+drgnpy_linux_helper_idr_for_each(PyObject *self, PyObject *args, PyObject *kwds)
+{
+	static char *keywords[] = {"idr", NULL};
+	struct drgn_error *err;
+	DrgnObject *idr;
+	if (!PyArg_ParseTupleAndKeywords(args, kwds, "O!:idr_for_each",
+					 keywords, &DrgnObject_type, &idr))
+		return NULL;
+
+	GenericIterator *iterator =
+		(GenericIterator *)GenericIterator_type.tp_alloc(
+			&GenericIterator_type, 0);
+	if (!iterator)
+		return NULL;
+	iterator->prog = DrgnObject_prog(idr);
+	Py_INCREF(iterator->prog);
+	iterator->next = idr_for_each_next;
+	iterator->iter_deinit = (void (*)(void *))linux_helper_idr_iter_deinit;
+	iterator->iter = malloc(sizeof(struct linux_helper_idr_iter));
+	if (!iterator->iter) {
+		Py_DECREF(iterator);
+		return (GenericIterator *)PyErr_NoMemory();
+	}
+	err = linux_helper_idr_iter_init(iterator->iter, &idr->obj);
+	if (err) {
+		Py_DECREF(iterator);
+		return set_drgn_error(err);
+	}
+	return iterator;
+}
+
+static PyObject *radix_tree_for_each_next(GenericIterator *self)
+{
+	struct linux_helper_radix_tree_iter_entry *entry;
+	struct drgn_error *err = linux_helper_radix_tree_iter_next(self->iter, &entry);
+	if (err)
+		return set_drgn_error(err);
+	if (!entry)
+		return NULL;
+	return idr_iter_entry_wrap(entry, self->prog);
+}
+
+GenericIterator *drgnpy_linux_helper_radix_tree_for_each(PyObject *self,
+							 PyObject *args,
+							 PyObject *kwds)
+{
+	static char *keywords[] = {"root", NULL};
+	struct drgn_error *err;
+	DrgnObject *root;
+	if (!PyArg_ParseTupleAndKeywords(args, kwds, "O!:radix_tree_for_each",
+					 keywords, &DrgnObject_type, &root))
+		return NULL;
+
+	GenericIterator *iterator =
+		(GenericIterator *)GenericIterator_type.tp_alloc(
+			&GenericIterator_type, 0);
+	if (!iterator)
+		return NULL;
+	iterator->prog = DrgnObject_prog(root);
+	Py_INCREF(iterator->prog);
+	iterator->next = radix_tree_for_each_next;
+	iterator->iter_deinit = (void (*)(void *))linux_helper_radix_tree_iter_deinit;
+	iterator->iter = malloc(sizeof(struct linux_helper_radix_tree_iter));
+	if (!iterator->iter) {
+		Py_DECREF(iterator);
+		return (GenericIterator *)PyErr_NoMemory();
+	}
+	err = linux_helper_radix_tree_iter_init(iterator->iter, &root->obj);
+	if (err) {
+		Py_DECREF(iterator);
+		return set_drgn_error(err);
+	}
+	return iterator;
+}

libdrgn/python/module.c

@@ -141,6 +141,18 @@ static PyMethodDef drgn_methods[] = {
 	{"_linux_helper_pgtable_l5_enabled",
 	 (PyCFunction)drgnpy_linux_helper_pgtable_l5_enabled,
 	 METH_VARARGS | METH_KEYWORDS},
+	{"_linux_helper_for_each_task",
+	 (PyCFunction)drgnpy_linux_helper_for_each_task,
+	 METH_VARARGS | METH_KEYWORDS, drgn__linux_helper_for_each_task_DOC},
+	{"_linux_helper_for_each_pid",
+	 (PyCFunction)drgnpy_linux_helper_for_each_pid,
+	 METH_VARARGS | METH_KEYWORDS, drgn__linux_helper_for_each_pid_DOC},
+	{"_linux_helper_idr_for_each",
+	 (PyCFunction)drgnpy_linux_helper_idr_for_each,
+	 METH_VARARGS | METH_KEYWORDS, drgn__linux_helper_idr_for_each_DOC},
+	{"_linux_helper_radix_tree_for_each",
+	 (PyCFunction)drgnpy_linux_helper_radix_tree_for_each,
+	 METH_VARARGS | METH_KEYWORDS, drgn__linux_helper_radix_tree_for_each_DOC},
 	{},
 };
 
@@ -230,6 +242,7 @@ DRGNPY_PUBLIC PyMODINIT_FUNC PyInit__drgn(void)
 	    add_type(m, &StackTrace_type) ||
 	    add_type(m, &Symbol_type) ||
 	    add_type(m, &DrgnType_type) ||
+	    add_type(m, &GenericIterator_type) ||
 	    add_type(m, &TypeEnumerator_type) ||
 	    add_type(m, &TypeMember_type) ||
 	    add_type(m, &TypeParameter_type) ||

tests/helpers/linux/test_pid.py

@@ -1,6 +1,7 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # SPDX-License-Identifier: GPL-3.0-or-later
 
+from multiprocessing import Barrier, Process
 import os
 
 from drgn.helpers.linux.pid import find_pid, find_task, for_each_pid, for_each_task
@@ -30,5 +31,23 @@ class TestPid(LinuxHelperTestCase):
         self.assertEqual(task.comm.string_(), comm)
 
     def test_for_each_task(self):
-        pid = os.getpid()
+        NUM_PROCS = 12
-        self.assertTrue(any(task.pid == pid for task in for_each_task(self.prog)))
+        barrier = Barrier(NUM_PROCS + 1)
+
+        def proc_func():
+            barrier.wait()
+
+        try:
+            procs = [Process(target=proc_func) for _ in range(NUM_PROCS)]
+            for proc in procs:
+                proc.start()
+            pids = {task.pid.value_() for task in for_each_task(self.prog)}
+            for proc in procs:
+                self.assertIn(proc.pid, pids)
+            self.assertIn(os.getpid(), pids)
+            barrier.wait()
+        except:
+            barrier.abort()
+            for proc in procs:
+                proc.terminate()
+            raise