Merge pull request #1203 from etsal/match-pthread-set-name

scx_layered: add rule for matching thread to name change

scheds/rust/scx_layered/src/bpf/intf.h

@@ -56,7 +56,14 @@ enum consts {
 	NSEC_PER_USEC		= 1000ULL,
 	NSEC_PER_MSEC		= (1000ULL * NSEC_PER_USEC),
 	MSEC_PER_SEC		= 1000ULL,
-	NSEC_PER_SEC		= NSEC_PER_MSEC * MSEC_PER_SEC
+	NSEC_PER_SEC		= NSEC_PER_MSEC * MSEC_PER_SEC,
+
+	SCXCMD_OP_NONE 		= 0,
+	SCXCMD_OP_JOIN 		= 1,
+	SCXCMD_OP_LEAVE 	= 2,
+
+	SCXCMD_PREFIX		= 0x5C10,
+	SCXCMD_COMLEN		= 13,
 };
 
 static inline void ___consts_sanity_check___(void) {
@@ -233,6 +240,7 @@ enum layer_match_kind {
 	MATCH_TGID_EQUALS,
 	MATCH_NSPID_EQUALS,
 	MATCH_NS_EQUALS,
+	MATCH_SCXCMD_JOIN,
 
 	NR_LAYER_MATCH_KINDS,
 };
@@ -307,4 +315,10 @@ struct layer {
 	char			name[MAX_LAYER_NAME];
 };
 
+struct scx_cmd {
+	u16			prefix;
+	u8 			opcode;
+	u8			cmd[SCXCMD_COMLEN];
+} __attribute__((packed));
+
 #endif /* __INTF_H */

scheds/rust/scx_layered/src/bpf/main.bpf.c

@@ -408,6 +408,8 @@ struct task_ctx {
 	/* for llcc->queue_runtime */
 	u32			qrt_layer_id;
 	u32			qrt_llc_id;
+
+	char 			join_layer[SCXCMD_COMLEN];
 };
 
 struct {
@@ -491,13 +493,59 @@ int BPF_PROG(tp_cgroup_attach_task, struct cgroup *cgrp, const char *cgrp_path,
 	return 0;
 }
 
+static int handle_cmd(struct task_ctx *taskc, struct scx_cmd *cmd)
+{
+
+	_Static_assert(sizeof(*cmd) == MAX_COMM, "scx_cmd has wrong size");
+
+	bpf_printk("received cmd");
+
+	/* Is this a valid command? */
+	if (cmd->prefix != SCXCMD_PREFIX)
+		return 0;
+
+	switch (cmd->opcode) {
+	case SCXCMD_OP_NONE:
+		break;
+
+	case SCXCMD_OP_JOIN:
+		__builtin_memcpy(taskc->join_layer, cmd->cmd, SCXCMD_COMLEN);
+		break;
+
+	case SCXCMD_OP_LEAVE:
+		__builtin_memset(taskc->join_layer, 0, SCXCMD_COMLEN);
+		break;
+
+	default:
+		break;
+	}
+
+	return 0;
+}
+
+
 SEC("tp_btf/task_rename")
 int BPF_PROG(tp_task_rename, struct task_struct *p, const char *buf)
 {
 	struct task_ctx *taskc;
+	struct scx_cmd cmd;
+	int ret;
+
+	if (!(taskc = lookup_task_ctx_may_fail(p))) {
+		bpf_printk("could not find task on rename");
+		return -EINVAL;
+	}
+
+	taskc->refresh_layer = true;
+
+	ret = bpf_probe_read_str(&cmd, sizeof(cmd), buf);
+	if (ret < 0) {
+		bpf_printk("could not new task name on rename");
+		return -EINVAL;
+	}
+
+	handle_cmd(taskc, &cmd);
 
-	if ((taskc = lookup_task_ctx_may_fail(p)))
-		taskc->refresh_layer = true;
 	return 0;
 }
 
@@ -1768,6 +1816,20 @@ static __noinline bool match_one(struct layer_match *match,
 		bpf_rcu_read_unlock();
 		return nsid == match->nsid;
 	}
+	case MATCH_SCXCMD_JOIN: {
+		struct task_ctx *taskc = lookup_task_ctx_may_fail(p);
+		if (!taskc) {
+			scx_bpf_error("could not find task");
+			return false;
+		}
+
+		/* The empty string means "no join command". */
+		if (!taskc->join_layer[0])
+			return false;
+
+		return match_prefix(match->comm_prefix, taskc->join_layer,
+			SCXCMD_COMLEN);
+	}
 	default:
 		scx_bpf_error("invalid match kind %d", match->kind);
 		return result;

scheds/rust/scx_layered/src/config.rs

@@ -73,6 +73,7 @@ pub enum LayerMatch {
     TGIDEquals(u32),
     NSPIDEquals(u64, u32),
     NSEquals(u32),
+    CmdJoin(String),
 }
 
 #[derive(Clone, Debug, Serialize, Deserialize)]

scheds/rust/scx_layered/src/main.rs

@@ -1170,6 +1170,10 @@ impl<'a> Scheduler<'a> {
                             mt.kind = bpf_intf::layer_match_kind_MATCH_NS_EQUALS as i32;
                             mt.nsid = *nsid as u64;
                         }
+                        LayerMatch::CmdJoin(joincmd) => {
+                            mt.kind = bpf_intf::consts_SCXCMD_OP_JOIN as i32;
+                            copy_into_cstr(&mut mt.comm_prefix, joincmd);
+                        }
                     }
                 }
                 layer.matches[or_i].nr_match_ands = or.len() as i32;