diff --git a/scheds/rust/scx_rusty/src/bpf/main.bpf.c b/scheds/rust/scx_rusty/src/bpf/main.bpf.c
index ea80841..ee95628 100644
--- a/scheds/rust/scx_rusty/src/bpf/main.bpf.c
+++ b/scheds/rust/scx_rusty/src/bpf/main.bpf.c
@@ -1133,10 +1133,28 @@ void BPF_STRUCT_OPS(rusty_set_cpumask, struct task_struct *p,
 			bpf_cpumask_subset((const struct cpumask *)all_cpumask, cpumask);
 }
 
+static s32 create_save_cpumask(struct bpf_cpumask **kptr)
+{
+	struct bpf_cpumask *cpumask;
+
+	cpumask = bpf_cpumask_create();
+	if (!cpumask) {
+		scx_bpf_error("Failed to create cpumask");
+		return -ENOMEM;
+	}
+
+	cpumask = bpf_kptr_xchg(kptr, cpumask);
+	if (cpumask) {
+		scx_bpf_error("kptr already had cpumask");
+		bpf_cpumask_release(cpumask);
+	}
+
+	return 0;
+}
+
 s32 BPF_STRUCT_OPS(rusty_init_task, struct task_struct *p,
 		   struct scx_init_task_args *args)
 {
-	struct bpf_cpumask *cpumask;
 	struct task_ctx taskc = { .dom_active_pids_gen = -1 };
 	struct task_ctx *map_value;
 	long ret;
@@ -1167,30 +1185,16 @@ s32 BPF_STRUCT_OPS(rusty_init_task, struct task_struct *p,
 		/* Should never happen -- it was just inserted above. */
 		return -EINVAL;
 
-	cpumask = bpf_cpumask_create();
-	if (!cpumask) {
+	ret = create_save_cpumask(&map_value->cpumask);
+	if (ret) {
 		bpf_map_delete_elem(&task_data, &pid);
-		return -ENOMEM;
+		return ret;
 	}
 
-	cpumask = bpf_kptr_xchg(&map_value->cpumask, cpumask);
-	if (cpumask) {
-		/* Should never happen as we just inserted it above. */
-		bpf_cpumask_release(cpumask);
+	ret = create_save_cpumask(&map_value->tmp_cpumask);
+	if (ret) {
 		bpf_map_delete_elem(&task_data, &pid);
-		return -EINVAL;
-	}
-
-	cpumask = bpf_cpumask_create();
-	if (!cpumask) {
-		scx_bpf_error("Failed to create BPF cpumask for task");
-		return -ENOMEM;
-	}
-	cpumask = bpf_kptr_xchg(&map_value->tmp_cpumask, cpumask);
-	if (cpumask) {
-		scx_bpf_error("%s[%d] tmp_cpumask already present", p->comm, p->pid);
-		bpf_cpumask_release(cpumask);
-		return -EEXIST;
+		return ret;
 	}
 
 	task_pick_and_set_domain(map_value, p, p->cpus_ptr, true);
@@ -1222,6 +1226,7 @@ static s32 create_node(u32 node_id)
 	u32 cpu;
 	struct bpf_cpumask *cpumask;
 	struct node_ctx *nodec;
+	s32 ret;
 
 	nodec = bpf_map_lookup_elem(&node_data, &node_id);
 	if (!nodec) {
@@ -1230,39 +1235,41 @@ static s32 create_node(u32 node_id)
 		return -ENOENT;
 	}
 
-	cpumask = bpf_cpumask_create();
-	if (!cpumask)
-		return -ENOMEM;
+	ret = create_save_cpumask(&nodec->cpumask);
+	if (ret)
+		return ret;
 
-	for (cpu = 0; cpu < MAX_CPUS; cpu++) {
+	bpf_rcu_read_lock();
+	cpumask = nodec->cpumask;
+	if (!cpumask) {
+		bpf_rcu_read_unlock();
+		scx_bpf_error("Failed to lookup node cpumask");
+		return -ENOENT;
+	}
+
+	bpf_for(cpu, 0, MAX_CPUS) {
 		const volatile u64 *nmask;
 
 		nmask = MEMBER_VPTR(numa_cpumasks, [node_id][cpu / 64]);
 		if (!nmask) {
 			scx_bpf_error("array index error");
-			bpf_cpumask_release(cpumask);
-			return -ENOENT;
+			ret = -ENOENT;
+			break;
 		}
 
 		if (*nmask & (1LLU << (cpu % 64)))
 			bpf_cpumask_set_cpu(cpu, cpumask);
 	}
 
-	cpumask = bpf_kptr_xchg(&nodec->cpumask, cpumask);
-	if (cpumask) {
-		scx_bpf_error("Node %u cpumask already present", node_id);
-		bpf_cpumask_release(cpumask);
-		return -EEXIST;
-	}
-
-	return 0;
+	bpf_rcu_read_unlock();
+	return ret;
 }
 
 static s32 create_dom(u32 dom_id)
 {
 	struct dom_ctx *domc;
 	struct node_ctx *nodec;
-	struct bpf_cpumask *cpumask, *node_mask;
+	struct bpf_cpumask *dom_mask, *node_mask, *all_mask;
 	u32 cpu, node_id;
 	s32 ret;
 
@@ -1286,52 +1293,42 @@ static s32 create_dom(u32 dom_id)
 		return -ENOENT;
 	}
 
-	cpumask = bpf_cpumask_create();
-	if (!cpumask) {
-		scx_bpf_error("Failed to create BPF cpumask for domain %u", dom_id);
-		return -ENOMEM;
+
+	ret = create_save_cpumask(&domc->cpumask);
+	if (ret)
+		return ret;
+
+	bpf_rcu_read_lock();
+	dom_mask = domc->cpumask;
+	all_mask = all_cpumask;
+	if (!dom_mask || !all_mask) {
+		bpf_rcu_read_unlock();
+		scx_bpf_error("Could not find cpumask");
+		return -ENOENT;
 	}
 
-	for (cpu = 0; cpu < MAX_CPUS; cpu++) {
+	bpf_for(cpu, 0, MAX_CPUS) {
 		const volatile u64 *dmask;
 
 		dmask = MEMBER_VPTR(dom_cpumasks, [dom_id][cpu / 64]);
 		if (!dmask) {
 			scx_bpf_error("array index error");
-			bpf_cpumask_release(cpumask);
-			return -ENOENT;
+			ret = -ENOENT;
+			break;
 		}
 
 		if (*dmask & (1LLU << (cpu % 64))) {
-			bpf_cpumask_set_cpu(cpu, cpumask);
-
-			bpf_rcu_read_lock();
-			if (all_cpumask)
-				bpf_cpumask_set_cpu(cpu, all_cpumask);
-			bpf_rcu_read_unlock();
+			bpf_cpumask_set_cpu(cpu, dom_mask);
+			bpf_cpumask_set_cpu(cpu, all_mask);
 		}
 	}
+	bpf_rcu_read_unlock();
+	if (ret)
+		return ret;
 
-	cpumask = bpf_kptr_xchg(&domc->cpumask, cpumask);
-	if (cpumask) {
-		scx_bpf_error("Domain %u cpumask already present", dom_id);
-		bpf_cpumask_release(cpumask);
-		return -EEXIST;
-	}
-
-	cpumask = bpf_cpumask_create();
-	if (!cpumask) {
-		scx_bpf_error("Failed to create BPF cpumask for domain %u",
-			      dom_id);
-		return -ENOMEM;
-	}
-	cpumask = bpf_kptr_xchg(&domc->direct_greedy_cpumask, cpumask);
-	if (cpumask) {
-		scx_bpf_error("Domain %u direct_greedy_cpumask already present",
-			      dom_id);
-		bpf_cpumask_release(cpumask);
-		return -EEXIST;
-	}
+	ret = create_save_cpumask(&domc->direct_greedy_cpumask);
+	if (ret)
+		return ret;
 
 	nodec = bpf_map_lookup_elem(&node_data, &node_id);
 	if (!nodec) {
@@ -1339,30 +1336,20 @@ static s32 create_dom(u32 dom_id)
 		scx_bpf_error("No node%u", node_id);
 		return -ENOENT;
 	}
+	ret = create_save_cpumask(&domc->node_cpumask);
+	if (ret)
+		return ret;
+
 	bpf_rcu_read_lock();
 	node_mask = nodec->cpumask;
-	if (!node_mask) {
+	dom_mask = domc->node_cpumask;
+	if (!node_mask || !dom_mask) {
 		bpf_rcu_read_unlock();
-		scx_bpf_error("NUMA %d mask not found for domain %u",
-			      node_id, dom_id);
+		scx_bpf_error("cpumask lookup failed");
 		return -ENOENT;
 	}
-	cpumask = bpf_cpumask_create();
-	if (!cpumask) {
-		bpf_rcu_read_unlock();
-		scx_bpf_error("Failed to create BPF cpumask for domain %u",
-			      dom_id);
-		return -ENOMEM;
-	}
-	bpf_cpumask_copy(cpumask, (const struct cpumask *)node_mask);
+	bpf_cpumask_copy(dom_mask, (const struct cpumask *)node_mask);
 	bpf_rcu_read_unlock();
-	cpumask = bpf_kptr_xchg(&domc->node_cpumask, cpumask);
-	if (cpumask) {
-		scx_bpf_error("Domain %u node_cpumask already present",
-			      dom_id);
-		bpf_cpumask_release(cpumask);
-		return -EEXIST;
-	}
 
 	return 0;
 }
@@ -1446,29 +1433,19 @@ void BPF_STRUCT_OPS(rusty_cpu_offline, s32 cpu)
 
 s32 BPF_STRUCT_OPS_SLEEPABLE(rusty_init)
 {
-	struct bpf_cpumask *cpumask;
 	s32 i, ret;
 
-	cpumask = bpf_cpumask_create();
-	if (!cpumask)
-		return -ENOMEM;
-	cpumask = bpf_kptr_xchg(&all_cpumask, cpumask);
-	if (cpumask)
-		bpf_cpumask_release(cpumask);
+	ret = create_save_cpumask(&all_cpumask);
+	if (ret)
+		return ret;
 
-	cpumask = bpf_cpumask_create();
-	if (!cpumask)
-		return -ENOMEM;
-	cpumask = bpf_kptr_xchg(&direct_greedy_cpumask, cpumask);
-	if (cpumask)
-		bpf_cpumask_release(cpumask);
+	ret = create_save_cpumask(&direct_greedy_cpumask);
+	if (ret)
+		return ret;
 
-	cpumask = bpf_cpumask_create();
-	if (!cpumask)
-		return -ENOMEM;
-	cpumask = bpf_kptr_xchg(&kick_greedy_cpumask, cpumask);
-	if (cpumask)
-		bpf_cpumask_release(cpumask);
+	ret = create_save_cpumask(&kick_greedy_cpumask);
+	if (ret)
+		return ret;
 
 	if (!switch_partial)
 		__COMPAT_scx_bpf_switch_all();