1#include <math/ilog2.h>
2#include <sch/sched.h>
3#include <stdatomic.h>
4#include <stdbool.h>
5#include <stddef.h>
6#include <stdint.h>
7#include <sync/spinlock.h>
8#include <thread/apc.h>
9
10#include "internal.h"
11#include "sched_profiling.h"
12
13bool scheduler_can_take_thread(size_t core, struct thread *target) {
14 if (thread_get_flags(t: target) & THREAD_FLAG_PINNED)
15 return false;
16
17 if (!cpu_mask_test(m: &target->allowed_cpus, cpu: core))
18 return false;
19
20 return atomic_load_explicit(&target->migrate_to, memory_order_acquire) ==
21 -1;
22}
23
24/* self->stealing_work should already be set before this is called */
25/* TODO: Rate limit me so I don't do a full scan of all cores due to that being
26 * expensive */
27struct scheduler *scheduler_pick_victim(struct scheduler *self) {
28 /* Ideally, we want to steal from our busiest core */
29 uint64_t max_thread_count = 0;
30 struct scheduler *victim = NULL;
31
32 size_t i;
33 for_each_cpu_id(i) {
34 struct scheduler *potential_victim = global.schedulers[i];
35
36 /* duh.... */
37 if (potential_victim == self)
38 continue;
39
40 bool victim_busy = atomic_load(&potential_victim->being_robbed) ||
41 atomic_load(&potential_victim->stealing_work);
42
43 size_t victim_thread_count = potential_victim->total_thread_count;
44 if (global.cores && global.cores[i] &&
45 !scheduler_core_idle(c: global.cores[i]))
46 victim_thread_count++;
47
48 uint64_t victim_scaled = victim_thread_count * 100;
49 uint64_t scaled =
50 self->total_thread_count * scheduler_data.steal_min_diff;
51 bool victim_is_poor = victim_scaled < scaled;
52
53 if (victim_busy || victim_is_poor)
54 continue;
55
56 if (potential_victim->total_thread_count > max_thread_count) {
57 max_thread_count = victim_thread_count;
58 victim = potential_victim;
59 }
60 }
61
62 if (victim)
63 atomic_store(&victim->being_robbed, true);
64
65 return victim;
66}
67
68static struct thread *steal_from_thread_rbt(struct scheduler *victim,
69 struct rbt *tree) {
70 struct rbt_node *node;
71 rbt_for_each_reverse(node, tree) {
72 struct thread *target = thread_from_rq_rbt_node(node);
73
74 /* we must first set the thread as `being_moved` before we
75 * check if we can steal the thread... */
76 if (!scheduler_can_take_thread(core: smp_core_id(), target))
77 continue;
78
79 rbt_delete(tree, z: node);
80
81 scheduler_decrement_thread_count(sched: victim, t: target);
82 return target;
83 }
84
85 /* Nothing found here */
86 return NULL;
87}
88
89static struct thread *steal_from_ts_threads(struct scheduler *victim) {
90 struct thread *stolen;
91
92 /* We first try to pick from threads that have not run this period */
93 stolen = steal_from_thread_rbt(victim, tree: &victim->thread_rbt);
94 if (stolen)
95 return stolen;
96
97 /* Nothing found? Let's try from the completed threads this period */
98 stolen = steal_from_thread_rbt(victim, tree: &victim->completed_rbt);
99 if (stolen)
100 return stolen;
101
102 return NULL;
103}
104
105static struct thread *steal_from_special_threads(struct scheduler *victim,
106 struct list_head *q) {
107 if (list_empty(head: q))
108 return NULL;
109
110 size_t core = smp_core_id();
111
112 struct list_head *pos, *n;
113 list_for_each_safe(pos, n, q) {
114 struct thread *t = thread_from_rq_list_node(pos);
115
116 if (!scheduler_can_take_thread(core, target: t)) {
117 continue;
118 }
119
120 kassert(thread_get_state(t) == THREAD_STATE_READY);
121
122 list_del_init(entry: &t->rq_list_node);
123
124 scheduler_decrement_thread_count(sched: victim, t);
125 return t;
126 }
127
128 return NULL;
129}
130
131struct thread *scheduler_steal_work(struct scheduler *new,
132 struct scheduler *victim) {
133 /* do not wait in a loop */
134 if (!spin_trylock_raw(lock: &victim->lock))
135 return NULL;
136
137 struct thread *stolen = NULL;
138 uint8_t mask = atomic_load(&victim->queue_bitmap);
139 while (mask) {
140 int level = 31 - __builtin_clz((uint32_t) mask);
141 mask &= ~(1ULL << level); /* remove that bit from local copy */
142
143 if (level == THREAD_PRIO_CLASS_TIMESHARE) {
144 stolen = steal_from_ts_threads(victim);
145 if (stolen)
146 break;
147
148 } else {
149 struct list_head *q =
150 scheduler_get_this_thread_queue(sched: victim, prio: level);
151
152 stolen = steal_from_special_threads(victim, q);
153 if (stolen)
154 break;
155 }
156 }
157
158 if (stolen) {
159 thread_set_runqueue(t: stolen, s: new);
160 thread_post_migrate(t: stolen, old_cpu: victim->core_id, new_cpu: new->core_id);
161 }
162
163 spin_unlock_raw(lock: &victim->lock);
164 return stolen;
165}
166
167static inline void begin_steal(struct scheduler *sched) {
168 atomic_store(&sched->stealing_work, true);
169}
170
171static inline bool try_begin_steal() {
172 unsigned current = atomic_load(&scheduler_data.active_stealers);
173 while (current < scheduler_data.max_concurrent_stealers) {
174 if (atomic_compare_exchange_weak(&scheduler_data.active_stealers,
175 &current, current + 1)) {
176 return true;
177 }
178 }
179 return false;
180}
181
182static inline void stop_steal(struct scheduler *sched,
183 struct scheduler *victim) {
184 if (victim)
185 atomic_store(&victim->being_robbed, false);
186
187 atomic_store(&sched->stealing_work, false);
188 atomic_fetch_sub(&scheduler_data.active_stealers, 1);
189}
190
191struct thread *scheduler_try_do_steal(struct scheduler *sched) {
192 if (!scheduler_can_steal_work(sched))
193 return NULL;
194
195 if (!try_begin_steal())
196 return NULL;
197
198 begin_steal(sched);
199 struct scheduler *victim = scheduler_pick_victim(self: sched);
200
201 if (!victim) {
202 stop_steal(sched, victim);
203 return NULL;
204 }
205
206 struct thread *stolen = scheduler_steal_work(new: sched, victim);
207 stop_steal(sched, victim);
208
209 if (stolen) {
210 sched_profiling_record_steal();
211 } else {
212 scheduler_try_push_to_idle_core(sched);
213 }
214
215 return stolen;
216}
217
218uint64_t scheduler_compute_steal_threshold() {
219 uint64_t threads = atomic_load(&scheduler_data.total_threads);
220 uint64_t threads_per_core = threads / global.core_count;
221
222 if (threads_per_core <= 1)
223 return 150;
224
225 if (threads_per_core >= 64)
226 return 110;
227
228 uint8_t log = ilog2(x: threads_per_core);
229 return 150 - (log * 5);
230}
231
232bool scheduler_can_steal_work(struct scheduler *sched) {
233 uint64_t val = atomic_load(&scheduler_data.total_threads);
234 uint64_t avg_core_threads = val / global.core_count;
235 return sched->total_thread_count < avg_core_threads;
236}
237