1#include <atomic.h>
2#include <compiler/intrinsic.h>
3#include <math/bit.h>
4#include <sch/sched.h>
5#include <stdbool.h>
6#include <stddef.h>
7#include <stdint.h>
8#include <sync/spinlock.h>
9#include <thread/apc.h>
10
11#include "internal.h"
12
13bool scheduler_can_take_thread(size_t core, struct thread *target) {
14 if (thread_is_pinned(t: target))
15 return false;
16
17 if (!cpu_mask_test(m: &target->allowed_cpus, cpu: core))
18 return false;
19
20 return atomic_load_acq(&target->migrate_to) == -1;
21}
22
23/* self->stealing_work should already be set before this is called */
24/* TODO: Rate limit me so I don't do a full scan of all cores due to that being
25 * expensive */
26struct scheduler *scheduler_pick_victim(struct scheduler *self) {
27 /* Ideally, we want to steal from our busiest core */
28 uint64_t max_thread_count = 0;
29 struct scheduler *victim = NULL;
30
31 size_t i;
32 for_each_cpu_id(i) {
33 struct scheduler *potential_victim = global.schedulers[i];
34
35 /* duh.... */
36 if (potential_victim == self)
37 continue;
38
39 bool victim_busy = atomic_load(&potential_victim->being_robbed) ||
40 atomic_load(&potential_victim->stealing_work);
41
42 size_t victim_thread_count = potential_victim->total_thread_count;
43 if (global.cores && global.cores[i] &&
44 !scheduler_core_idle(c: global.cores[i]))
45 victim_thread_count++;
46
47 uint64_t victim_scaled = victim_thread_count * 100;
48 uint64_t scaled =
49 (uint64_t) self->total_thread_count *
50 (uint64_t) atomic_load_rlx(&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_id(cond: TOPC_IRQL), 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_head_empty(h: q))
108 return NULL;
109
110 size_t core = smp_id(cond: TOPC_IRQL);
111
112 struct list_node *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 enforce(thread_get_state(t) == THREAD_STATE_READY);
121
122 list_node_delete_init(n: &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(&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 - ci_clz((uint32_t) mask);
141 mask = BIT_CLEAR(mask, 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 scheduler_debug_validate(sched: victim);
162 }
163
164 spin_unlock_raw(&victim->lock);
165 return stolen;
166}
167
168static inline void begin_steal(struct scheduler *sched) {
169 atomic_store(&sched->stealing_work, true);
170}
171
172static inline bool try_begin_steal(void) {
173 uint32_t current = atomic_load_rlx(&scheduler_data.active_stealers);
174 while (current < scheduler_data.max_concurrent_stealers) {
175 if (atomic_cas_weak(&scheduler_data.active_stealers, &current,
176 current + 1, mo_acquire, mo_rlx)) {
177 return true;
178 }
179 }
180 return false;
181}
182
183static inline void stop_steal(struct scheduler *sched,
184 struct scheduler *victim) {
185 if (victim)
186 atomic_store(&victim->being_robbed, false);
187
188 atomic_store(&sched->stealing_work, false);
189 atomic_dec_release(&scheduler_data.active_stealers);
190}
191
192struct thread *scheduler_try_do_steal(struct scheduler *sched) {
193 if (!scheduler_can_steal_work(sched))
194 return NULL;
195
196 if (!try_begin_steal())
197 return NULL;
198
199 begin_steal(sched);
200 struct scheduler *victim = scheduler_pick_victim(self: sched);
201
202 if (!victim) {
203 stop_steal(sched, victim);
204 return NULL;
205 }
206
207 struct thread *stolen = scheduler_steal_work(new: sched, victim);
208 stop_steal(sched, victim);
209
210 if (!stolen)
211 scheduler_try_push_to_idle_core(sched);
212
213 return stolen;
214}
215
216uint64_t scheduler_compute_steal_threshold(void) {
217 uint64_t threads = atomic_load(&scheduler_data.total_threads);
218 uint64_t threads_per_core = threads / global.core_count;
219
220 if (threads_per_core <= 1)
221 return 150;
222
223 if (threads_per_core >= 64)
224 return 110;
225
226 uint8_t log = ilog2(x: threads_per_core);
227 return 150 - (log * 5);
228}
229
230bool scheduler_can_steal_work(struct scheduler *sched) {
231 uint64_t val = atomic_load(&scheduler_data.total_threads);
232 uint64_t avg_core_threads = val / global.core_count;
233
234 if (sched->total_thread_count == 0)
235 return val > 0;
236
237 return sched->total_thread_count < avg_core_threads;
238}
239