chaos: reimplement better random support

c576237 added random support to scx_chaos with less bias (very very
nearly 0). It was reverted because it broke random delays.

It turns out the random implementation was fine but the callsite was
wrong. Instead of adding the random delay to the current time it was
setting the target time to the random delay, which was always in the
past, and hence scheduling things immediately. With that fixed, this
appears to work.

Test plan:
- CI

```
 # wakeup latencies with no flags
 56.166µs |▁▄▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇███████████████ 132.125µs ███▇▇▇▇▆▆▆▃| 141.834µs
63.75µs |▁▄▄▄▄▅▅▅▅▆▆▆▆▆▆▆▆▆▆▆▆▇▇▇▇▇▇▇▇▇▇▇▇▇██████ 74.767µs ███████▇▇▇▇▇▇▇▇▇▇▇▇▇▆▆▆▆▆▆▆▆▆▆▆▆▅▅▅▅▅▅▅▅▅▄▄▄▄▄▄▄▄▄▃▃| 88.6µs
65.835µs |▁▆▇▇ 84.235µs █████████████▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▅▅▅▅▅▅▅▅▅▅▅▄▄▄▄▄▄▄▄▄▄▃▃| 446.137µs
59.289µs |▁▅▆▆▆▆▆▇▇▇▇▇██ 117.979µs ████████████▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▅▅▅▅▅▅▅▄▄▄▄▄▄▄▄▃▃▂▁| 407.666µs
64.987µs |▁▄▅▅▆▆▆▆▇▇▇▇▇██ 148.62µs █████████████▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▅▅▅▄▄▄▃▃▃▃▂▂▂▂▁▁| 536.936µs
52.78µs |▁▂▂▃▃▄▄▅▆▆▆▆▆▆▇▇▇▇▇▇██ 146.406µs ██████████████▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▅▄▄▃▁| 415.672µs
58.819µs |▁▃▄▄▅▅▆▆▆▆▆▆▆▇▇▇▇▇▇▇▇▇████ 174.184µs █████████████▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▅▄▄▁| 437.196µs
62.938µs |▁▆▇ 127.197µs ██▇▇▇▇▇▆▆▆▅▄▃▃▃▃▃▃▃▃▃▃▃▃▃▃▃▃▃▃▃▃▃▃▃▃▃▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁| 2.105205ms
59.026µs |▁▂▄▄▅▆▆▆▆▆▆▆▆▆▆▆▆▇▇▇▇▇▇▇▇▇▇▇▇▇██████ 227.526µs ██████████▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▅▅▅▄▄▃▂▁| 456.578µs
Benchmarking sleep_wakeup_histogram/wakeup_latency: Collecting 10 samples in estimated 5.0415 s (490 iterations)
64.466µs |▁▃▄▄▄▅▅▅▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇██████████ 312.808µs ██████▇▇▇▇▇▇▇▇▇▆▆▆▆▆▆▆▆▆▅▅▄▄▁| 431.229µs
48.977µs |▁▂▄▄▅▅▆▆▆▆▆▇▇▇▇▇██ 125.064µs ██████████████▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▅▅▅▄▄▄▁| 410.671µs
58.827µs |▁▃▄▅▆▆▆▆▆▇▇▇▇▇██ 127.664µs ██████████████▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▅▅▅▄▄▄▁| 421.334µs
63.12µs |▁▃▄▄▅▅▆▆▆▆▇▇▇▇▇███ 135.895µs ██████████████▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▅▄▁| 424.716µs
37.173µs |▁▁▂▂▃▃▄▄▅▆▆▆▆▇▇▇▇██ 132.831µs ███████████▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▅▅▅▅▅▅▄▄▄▄▄▃▃▃▂▂▁| 470.139µs
60.392µs |▁▂▂▃▃▄▄▄▅▅▅▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇██████████████ 410.584µs ██▇▇▆▆▅▅▄▄▁| 457.748µs
58.813µs |▁▃▄▅▆▆▆▆▇▇▇▇▇█ 124.729µs ███████████████▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▅▅▄▄▁| 453.846µs
55.203µs |▁▃▄▄▅▅▆▆▆▆▇▇▇▇██ 130.237µs ██████████████▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▅▅▄▄▄▁| 466.242µs
62.011µs |▁▄▄▅▅▆▆▆▆▆▆▆▇▇▇▇▇▇▇▇▇███ 186.162µs ████████████▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▅▅▄▄▃▃▃▂▂▂▁▁| 499.317µs
44.48µs |▁▃▄▄▄▄▅▅▅▆▆▆▆▆▆▇▇▇▇▇▇██ 151.382µs █████████████▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▅▅▄▄▄▁| 450.172µs

 # wakeup latencies with --random-delay-frequency 1.0 --random-delay-min-us 100000 --random-delay-max-us 200000
67.984µs | 103.273µs ▁████████████▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▅▅▅▅▅▅▅▅▅▅▅▄▄▄▄▄▄▄▄▄▄▄▃▃| 149.695558ms
69.85µs |▁▆▆▆▆▆▆▆▆▆▆▆▇▇▇▇▇▇▇▇▇▇▇▇▇██████ 73.03691ms ████████████▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▅▄▃| 197.940289ms
66.854µs |▁▆▆▆▆▆▆▆▆▆▆▆▇▇▇▇▇▇▇▇▇▇▇▇▇▇███████ 79.471527ms ██████████▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▅▅▅▄▄▄▃| 200.842494ms
72.268µs |▁▆▆▆▆▆▆▆▆▆▆▆▆▇▇▇▇▇▇▇▇▇▇▇▇▇▇███████ 79.757623ms ██████████▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▅▅▄▄▁| 195.335728ms
45.997µs |▁▆▆▆▆▆▆▆▆▇▇▇▇▇▇▇▇▇▇▇█████ 59.430492ms ███████████▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▅▅▄▄▄▁| 201.135123ms
Benchmarking sleep_wakeup_histogram/wakeup_latency: Collecting 10 samples in estimated 6.1568 s (40 iterations)
72.645µs |▁▆▆▆▆▆▆▆▆▆▆▆▆▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇███████ 83.468765ms ████████▇▇▇▇▇▇▇▇▇▇▇▇▇▇▆▆▆▆▆▆▆▆▆▆▆▆▆▅▅▅▅▅▅▄▄▄▄▄▄▃▃| 198.469905ms
77.531µs |▁▆▆▆▆▆▆▆▆▆▆▆▆▆▆▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇████████ 92.318916ms █████████▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▅▄▃| 194.16041ms
67.939µs |▁▆▆▆▆▆▆▆▆▆▆▆▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇██████ 75.667921ms ██████████▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▅▅▅▄▄▃| 188.527733ms
66.176µs |▁▆▆▆▆▆▆▆▆▆▆▆▇▇▇▇▇▇▇▇▇▇▇▇▇▇██████ 75.595375ms ██████████▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▅▅▅▅▅▄▄▄▄▃| 193.914143ms
67.455µs |▁▆▆▆▆▆▆▆▆▆▆▆▇▇▇▇▇▇▇▇▇▇▇▇▇██████ 75.774696ms ██████████▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▅▅▅▅▅▄▄▄▄▃| 199.696546ms
61.654µs |▁▆▆▆▆▆▆▆▆▆▆▆▆▆▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇███████ 86.050604ms █████████▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▆▆▆▆▆▆▆▆▆▆▆▆▆▆▅▅▄▄▃| 186.936313ms
65.914µs |▁▆▆▆▆▆▆▆▆▆▆▆▇▇▇▇▇▇▇▇▇▇▇▇▇▇██████ 76.277914ms ██████████▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▅▅▅▅▄▄▄▃| 196.387308ms
91.47µs |▁▆▆▆▆▆▆▆▆▆▆▆▇▇▇▇▇▇▇▇▇▇▇▇▇▇██████ 76.357756ms ██████████▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▅▅▅▅▄▄▄▄▃| 200.616152ms
79.914µs |▁▆▆▆▆▆▆▆▆▆▆▆▆▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇███████ 83.185347ms ████████▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▆▆▆▆▆▆▆▆▆▆▆▆▆▆▅▅▅▅▅▄▄▄▄▄▄▃| 198.712172ms
77.44µs |▁▆▆▆▆▆▆▆▆▆▆▇▇▇▇▇▇▇▇▇▇▇▇▇██████ 72.513261ms ████████████▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▅▄▄▃| 199.976316ms
```

Author: JakeHillion

scheds/include/scx/common.bpf.h

@@ -763,6 +763,14 @@ static inline u64 scale_by_task_weight_inverse(const struct task_struct *p, u64
 	return value * 100 / p->scx.weight;
 }
 
+/*
+ * Get a random u64 from the kernel's pseudo-random generator.
+ */
+static inline u64 get_prandom_u64()
+{
+	return ((u64)bpf_get_prandom_u32() << 32) | bpf_get_prandom_u32();
+}
+
 
 #include "compat.bpf.h"
 #include "enums.bpf.h"

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

@@ -10,10 +10,12 @@ typedef unsigned long long u64;
 #endif
 
 enum chaos_consts {
-	CHAOS_DSQ_BASE_SHIFT  = 16,
-	CHAOS_DSQ_BASE	      = 1 << CHAOS_DSQ_BASE_SHIFT,
+	CHAOS_DSQ_BASE_SHIFT	= 16,
+	CHAOS_DSQ_BASE		= 1 << CHAOS_DSQ_BASE_SHIFT,
 
-	CHAOS_NUM_PPIDS_CHECK = 1 << 20,
+	CHAOS_NUM_PPIDS_CHECK	= 1 << 20,
+
+	CHAOS_MAX_RAND_ATTEMPTS = 512,
 };
 
 enum chaos_match {

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

@@ -95,6 +95,47 @@ struct chaos_task_ctx *lookup_create_chaos_task_ctx(struct task_struct *p)
 				    BPF_LOCAL_STORAGE_GET_F_CREATE);
 }
 
+static __always_inline u64 chaos_get_prandom_u64_limit(u64 s)
+{
+	// Implementation of Lemire's algorithm 5 without 128-bit arithmetic.
+	// See https://arxiv.org/pdf/1805.10941v2 for details.
+	// Uses a bounded loop given this is BPF, but given the loop should
+	// rarely be entered this is fine.
+	u64 x, m_low, m_high;
+	u64 t;
+
+	x = get_prandom_u64();
+
+	// Compute 64-bit multiplication high and low parts
+	// m = x * s, split into m_high and m_low
+	m_high = ((x >> 32) * (s >> 32)) +
+		 (((x & 0xFFFFFFFF) * (s >> 32)) >> 32) +
+		 (((x >> 32) * (s & 0xFFFFFFFF)) >> 32);
+	m_low = x * s;
+
+	if (m_low < s) {
+		t = ((u64)(-s)) % s;
+		bpf_repeat(CHAOS_MAX_RAND_ATTEMPTS)
+		{
+			if (m_low >= t)
+				break;
+
+			x      = get_prandom_u64();
+			m_high = ((x >> 32) * (s >> 32)) +
+				 (((x & 0xFFFFFFFF) * (s >> 32)) >> 32) +
+				 (((x >> 32) * (s & 0xFFFFFFFF)) >> 32);
+			m_low = x * s;
+		}
+	}
+
+	return m_high;
+}
+
+static __always_inline u64 chaos_get_uniform_u64(u64 min, u64 max)
+{
+	return min + chaos_get_prandom_u64_limit(max - min + 1);
+}
+
 static __always_inline void chaos_stat_inc(enum chaos_stat_idx stat)
 {
 	u64 *cnt_p = bpf_map_lookup_elem(&chaos_stats, &stat);
@@ -270,15 +311,8 @@ __weak s32 enqueue_random_delay(struct task_struct *p	     __arg_trusted,
 				struct chaos_task_ctx *taskc __arg_nonnull,
 				u64 min_ns, u64 max_ns)
 {
-	u64 rand64 = ((u64)bpf_get_prandom_u32() << 32) | bpf_get_prandom_u32();
-
-	u64 vtime  = bpf_ktime_get_ns() + min_ns;
-	if (min_ns != max_ns) {
-		vtime += rand64 % (max_ns - min_ns);
-	}
-
+	u64 vtime = bpf_ktime_get_ns() + chaos_get_uniform_u64(min_ns, max_ns);
 	scx_bpf_dsq_insert_vtime(p, get_cpu_delay_dsq(-1), 0, vtime, enq_flags);
-
 	return true;
 }