Merge pull request #1103 from hodgesds/layered-ns-pid-target

scx_layered: Add pid namespace layer matching
sched-ext · Dec 13, 2024 · e8e68e8 · e8e68e8
2 parents 5333d25 + 2b49811
commit e8e68e8
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 0 deletions.
diff --git a/scheds/rust/scx_layered/src/bpf/intf.h b/scheds/rust/scx_layered/src/bpf/intf.h
@@ -220,6 +220,7 @@ enum layer_match_kind {
 	MATCH_PID_EQUALS,
 	MATCH_PPID_EQUALS,
 	MATCH_TGID_EQUALS,
+	MATCH_NSPID_EQUALS,
 
 	NR_LAYER_MATCH_KINDS,
 };
@@ -235,6 +236,7 @@ struct layer_match {
 	u32		pid;
 	u32		ppid;
 	u32		tgid;
+	u64		nsid;
 };
 
 struct layer_match_ands {

diff --git a/scheds/rust/scx_layered/src/bpf/main.bpf.c b/scheds/rust/scx_layered/src/bpf/main.bpf.c
@@ -5,8 +5,10 @@
 #endif
 #define LSP_INC
 #include "../../../../include/scx/common.bpf.h"
+#include "../../../../include/scx/namespace_impl.bpf.h"
 #else
 #include <scx/common.bpf.h>
+#include <scx/namespace_impl.bpf.h>
 #endif
 
 #include <errno.h>
@@ -1684,6 +1686,20 @@ static __noinline bool match_one(struct layer_match *match,
 		return p->real_parent->pid == match->ppid;
 	case MATCH_TGID_EQUALS:
 		return p->tgid == match->tgid;
+	case MATCH_NSPID_EQUALS:
+		// To do namespace pid matching we need to translate the root
+		// pid from bpf side to the namespace pid.
+		bpf_rcu_read_lock();
+		struct pid *p_pid = get_task_pid_ptr(p, PIDTYPE_PID);
+		struct pid_namespace *pid_ns = get_task_pid_ns(p, PIDTYPE_TGID);
+		if (!p_pid || !pid_ns) {
+			bpf_rcu_read_unlock();
+			return result;
+		}
+		pid_t nspid = get_pid_nr_ns(p_pid, pid_ns);
+		u64 nsid = BPF_CORE_READ(pid_ns, ns.inum);
+		bpf_rcu_read_unlock();
+		return (u32)nspid == match->pid && nsid == match->nsid;
 	default:
 		scx_bpf_error("invalid match kind %d", match->kind);
 		return result;
@@ -2649,6 +2665,10 @@ static s32 init_layer(int layer_id)
 			case MATCH_TGID_EQUALS:
 				dbg("%s TGID %u", header, match->tgid);
 				break;
+			case MATCH_NSPID_EQUALS:
+				dbg("%s NSID %lld PID %d",
+				    header, match->nsid, match->pid);
+				break;
 			default:
 				scx_bpf_error("%s Invalid kind", header);
 				return -EINVAL;

diff --git a/scheds/rust/scx_layered/src/config.rs b/scheds/rust/scx_layered/src/config.rs
@@ -71,6 +71,7 @@ pub enum LayerMatch {
     PIDEquals(u32),
     PPIDEquals(u32),
     TGIDEquals(u32),
+    NSPIDEquals(u64, u32),
 }
 
 #[derive(Clone, Debug, Serialize, Deserialize)]

diff --git a/scheds/rust/scx_layered/src/main.rs b/scheds/rust/scx_layered/src/main.rs
@@ -267,6 +267,8 @@ lazy_static! {
 ///
 /// - TGIDEquals: Matches if the task's tgid matches the value.
 ///
+/// - NSPIDEquals: Matches if the task's namespace id and pid matches the values.
+///
 /// While there are complexity limitations as the matches are performed in
 /// BPF, it is straightforward to add more types of matches.
 ///
@@ -1144,6 +1146,11 @@ impl<'a> Scheduler<'a> {
                             mt.kind = bpf_intf::layer_match_kind_MATCH_TGID_EQUALS as i32;
                             mt.tgid = *tgid;
                         }
+                        LayerMatch::NSPIDEquals(nsid, pid) => {
+                            mt.kind = bpf_intf::layer_match_kind_MATCH_NSPID_EQUALS as i32;
+                            mt.nsid = *nsid;
+                            mt.pid = *pid;
+                        }
                     }
                 }
                 layer.matches[or_i].nr_match_ands = or.len() as i32;