path: root/tools/sched_ext/scx_flatcg.bpf.c

/* SPDX-License-Identifier: GPL-2.0 */
/*
 * A demo sched_ext flattened cgroup hierarchy scheduler. It implements
 * hierarchical weight-based cgroup CPU control by flattening the cgroup
 * hierarchy into a single layer by compounding the active weight share at each
 * level. Consider the following hierarchy with weights in parentheses:
 *
 * R + A (100) + B (100)
 *   |         \ C (100)
 *   \ D (200)
 *
 * Ignoring the root and threaded cgroups, only B, C and D can contain tasks.
 * Let's say all three have runnable tasks. The total share that each of these
 * three cgroups is entitled to can be calculated by compounding its share at
 * each level.
 *
 * For example, B is competing against C and in that competition its share is
 * 100/(100+100) == 1/2. At its parent level, A is competing against D and A's
 * share in that competition is 100/(200+100) == 1/3. B's eventual share in the
 * system can be calculated by multiplying the two shares, 1/2 * 1/3 == 1/6. C's
 * eventual shaer is the same at 1/6. D is only competing at the top level and
 * its share is 200/(100+200) == 2/3.
 *
 * So, instead of hierarchically scheduling level-by-level, we can consider it
 * as B, C and D competing each other with respective share of 1/6, 1/6 and 2/3
 * and keep updating the eventual shares as the cgroups' runnable states change.
 *
 * This flattening of hierarchy can bring a substantial performance gain when
 * the cgroup hierarchy is nested multiple levels. in a simple benchmark using
 * wrk[8] on apache serving a CGI script calculating sha1sum of a small file, it
 * outperforms CFS by ~3% with CPU controller disabled and by ~10% with two
 * apache instances competing with 2:1 weight ratio nested four level deep.
 *
 * However, the gain comes at the cost of not being able to properly handle
 * thundering herd of cgroups. For example, if many cgroups which are nested
 * behind a low priority parent cgroup wake up around the same time, they may be
 * able to consume more CPU cycles than they are entitled to. In many use cases,
 * this isn't a real concern especially given the performance gain. Also, there
 * are ways to mitigate the problem further by e.g. introducing an extra
 * scheduling layer on cgroup delegation boundaries.
 *
 * The scheduler first picks the cgroup to run and then schedule the tasks
 * within by using nested weighted vtime scheduling by default. The
 * cgroup-internal scheduling can be switched to FIFO with the -f option.
 */
#include <scx/common.bpf.h>
#include "scx_flatcg.h"

/*
 * Maximum amount of retries to find a valid cgroup.
 */
enum {
	FALLBACK_DSQ		= 0,
	CGROUP_MAX_RETRIES	= 1024,
};

char _license[] SEC("license") = "GPL";

const volatile u32 nr_cpus = 32;	/* !0 for veristat, set during init */
const volatile u64 cgrp_slice_ns = SCX_SLICE_DFL;
const volatile bool fifo_sched;

u64 cvtime_now;
UEI_DEFINE(uei);

struct {
	__uint(type, BPF_MAP_TYPE_PERCPU_ARRAY);
	__type(key, u32);
	__type(value, u64);
	__uint(max_entries, FCG_NR_STATS);
} stats SEC(".maps");

static void stat_inc(enum fcg_stat_idx idx)
{
	u32 idx_v = idx;

	u64 *cnt_p = bpf_map_lookup_elem(&stats, &idx_v);
	if (cnt_p)
		(*cnt_p)++;
}

struct fcg_cpu_ctx {
	u64			cur_cgid;
	u64			cur_at;
};

struct {
	__uint(type, BPF_MAP_TYPE_PERCPU_ARRAY);
	__type(key, u32);
	__type(value, struct fcg_cpu_ctx);
	__uint(max_entries, 1);
} cpu_ctx SEC(".maps");

struct {
	__uint(type, BPF_MAP_TYPE_CGRP_STORAGE);
	__uint(map_flags, BPF_F_NO_PREALLOC);
	__type(key, int);
	__type(value, struct fcg_cgrp_ctx);
} cgrp_ctx SEC(".maps");

struct cgv_node {
	struct bpf_rb_node	rb_node;
	__u64			cvtime;
	__u64			cgid;
};

private(CGV_TREE) struct bpf_spin_lock cgv_tree_lock;
private(CGV_TREE) struct bpf_rb_root cgv_tree __contains(cgv_node, rb_node);

struct cgv_node_stash {
	struct cgv_node __kptr *node;
};

struct {
	__uint(type, BPF_MAP_TYPE_HASH);
	__uint(max_entries, 16384);
	__type(key, __u64);
	__type(value, struct cgv_node_stash);
} cgv_node_stash SEC(".maps");

struct fcg_task_ctx {
	u64		bypassed_at;
};

struct {
	__uint(type, BPF_MAP_TYPE_TASK_STORAGE);
	__uint(map_flags, BPF_F_NO_PREALLOC);
	__type(key, int);
	__type(value, struct fcg_task_ctx);
} task_ctx SEC(".maps");

/* gets inc'd on weight tree changes to expire the cached hweights */
u64 hweight_gen = 1;

static u64 div_round_up(u64 dividend, u64 divisor)
{
	return (dividend + divisor - 1) / divisor;
}

static bool vtime_before(u64 a, u64 b)
{
	return (s64)(a - b) < 0;
}

static bool cgv_node_less(struct bpf_rb_node *a, const struct bpf_rb_node *b)
{
	struct cgv_node *cgc_a, *cgc_b;

	cgc_a = container_of(a, struct cgv_node, rb_node);
	cgc_b = container_of(b, struct cgv_node, rb_node);

	return cgc_a->cvtime < cgc_b->cvtime;
}

static struct fcg_cpu_ctx *find_cpu_ctx(void)
{
	struct fcg_cpu_ctx *cpuc;
	u32 idx = 0;

	cpuc = bpf_map_lookup_elem(&cpu_ctx, &idx);
	if (!cpuc) {
		scx_bpf_error("cpu_ctx lookup failed");
		return NULL;
	}
	return cpuc;
}

static struct fcg_cgrp_ctx *find_cgrp_ctx(struct cgroup *cgrp)
{
	struct fcg_cgrp_ctx *cgc;

	cgc = bpf_cgrp_storage_get(&cgrp_ctx, cgrp, 0, 0);
	if (!cgc) {
		scx_bpf_error("cgrp_ctx lookup failed for cgid %llu", cgrp->kn->id);
		return NULL;
	}
	return cgc;
}

static struct fcg_cgrp_ctx *find_ancestor_cgrp_ctx(struct cgroup *cgrp, int level)
{
	struct fcg_cgrp_ctx *cgc;

	cgrp = bpf_cgroup_ancestor(cgrp, level);
	if (!cgrp) {
		scx_bpf_error("ancestor cgroup lookup failed");
		return NULL;
	}

	cgc = find_cgrp_ctx(cgrp);
	if (!cgc)
		scx_bpf_error("ancestor cgrp_ctx lookup failed");
	bpf_cgroup_release(cgrp);
	return cgc;
}

static void cgrp_refresh_hweight(struct cgroup *cgrp, struct fcg_cgrp_ctx *cgc)
{
	int level;

	if (!cgc->nr_active) {
		stat_inc(FCG_STAT_HWT_SKIP);
		return;
	}

	if (cgc->hweight_gen == hweight_gen) {
		stat_inc(FCG_STAT_HWT_CACHE);
		return;
	}

	stat_inc(FCG_STAT_HWT_UPDATES);
	bpf_for(level, 0, cgrp->level + 1) {
		struct fcg_cgrp_ctx *cgc;
		bool is_active;

		cgc = find_ancestor_cgrp_ctx(cgrp, level);
		if (!cgc)
			break;

		if (!level) {
			cgc->hweight = FCG_HWEIGHT_ONE;
			cgc->hweight_gen = hweight_gen;
		} else {
			struct fcg_cgrp_ctx *pcgc;

			pcgc = find_ancestor_cgrp_ctx(cgrp, level - 1);
			if (!pcgc)
				break;

			/*
			 * We can be opportunistic here and not grab the
			 * cgv_tree_lock and deal with the occasional races.
			 * However, hweight updates are already cached and
			 * r