RustForge

rust-kernel

Rustbelt: a hobby x86_64 kernel written in Rust

rust-kernel / src / thread.rs
31754 bytes ยท raw
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
//! Threads and the scheduler: preemptive, multi-core, with sleep/wake.
//!
//! Each thread has its own kernel stack (64 KiB with a guard page, see
//! stack.rs). Every CPU's APIC timer fires ~100 times a second and calls
//! `schedule()`, which puts the running thread at the back of the shared
//! **ready queue** and switches to the thread at the front. Threads don't
//! have to cooperate: even `loop {}` gets interrupted.
//!
//! How a switch works (inside the timer interrupt, interrupts off):
//!
//!   thread A running -> timer IRQ -> CPU pushes A's RIP/RSP/RFLAGS (the
//!   interrupt frame) on A's stack -> the `x86-interrupt` handler saves the
//!   caller-saved registers (rax, rcx, rdx, rsi, rdi, r8-r11) on A's stack ->
//!   `schedule()` -> `switch_context(&A.rsp, B.rsp)`:
//!       push the callee-saved registers (rbx, rbp, r12-r15) on A's stack
//!       save rsp into A.rsp
//!       load rsp from B.rsp            <- now we're on B's stack
//!       pop B's callee-saved registers
//!       ret                            <- returns into B's earlier schedule()
//!   ... which returns into B's timer handler, which restores B's caller-saved
//!   registers and `iretq`s back to wherever B was interrupted.
//!
//! So a waiting thread's whole CPU state is simply sitting on its own stack,
//! and all we have to remember per thread is its stack pointer.
//!
//! Threads that run user code (processes, see process.rs) also carry an
//! address space (CR3). On every switch the scheduler loads the next
//! thread's CR3 and points this CPU's TSS.RSP0 and syscall stack at the
//! thread's kernel stack.
//!
//! **States.** A thread is Ready (in the queue), Running (on some CPU),
//! Blocked (waiting on a `WaitQueue`: keyboard input, a pipe, a child...),
//! Sleeping (until a PIT tick; the PIT handler wakes it) or Finished. Blocked
//! and sleeping threads are not in the queue at all, so they cost no CPU
//! time: no busy-polling anywhere.
//!
//! **Idle threads.** Each CPU has one, run when the queue is empty. It
//! halts the CPU (`hlt`) until an interrupt arrives. A CPU that queues work
//! while another is idle sends it a "reschedule" IPI (smp.rs).
//!
//! **Several CPUs, one queue.** The scheduler's data is behind one spinlock.
//! Two subtle points:
//!   * A thread that was just preempted on CPU A is put in the queue *before*
//!     A has finished switching away from its stack. CPU B must not start
//!     running it on the same stack at the same time! So each thread has an
//!     `on_cpu` flag, cleared only after the switch (`finish_switch`), and
//!     the scheduler skips threads that still have it set.
//!   * Likewise a finished thread's stack (and maybe its process's memory)
//!     can only be freed once no CPU is on it: also in `finish_switch`.
//!
//! **Locking rule.** Any lock that is ever taken with interrupts *off* (the
//! scheduler's, a wait queue's, a pipe buffer's, the frame allocator's...)
//! must *always* be taken with interrupts off. Otherwise a thread could be
//! preempted while holding it, and an interrupt handler on the same CPU
//! that wants it would spin forever: the holder can't run until the
//! handler returns. (Spinning briefly while *another* CPU holds a lock is
//! fine: that CPU keeps running and will release it.)
//!
//! A *new* thread has never been switched away from, so `spawn` builds a
//! fake stack that looks like `switch_context` saved it, with the "return
//! address" pointing at `thread_trampoline`, which calls the thread's function.

use crate::errno::Errno;
use crate::process::{Process, Task};
use crate::smp::{self, MAX_CPUS};
use crate::stack::KernelStack;
use alloc::{
    boxed::Box,
    collections::{BTreeMap, VecDeque},
    string::String,
    sync::Arc,
    vec::Vec,
};
use core::arch::global_asm;
use core::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use spin::Mutex;
use x86_64::instructions::interrupts;
use x86_64::registers::control::Cr3;
use x86_64::structures::paging::PhysFrame;

/// Thread id. Never reused.
pub type Tid = usize;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum State {
    Ready,
    Running,
    Blocked,
    /// Sleeping until this PIT tick.
    Sleeping(u64),
    Finished,
}

struct Thread {
    name: String,
    /// Saved stack pointer while the thread is not running.
    rsp: u64,
    /// The thread's stack. `None` for the boot thread "main", which runs on
    /// the stack the bootloader gave us (that one has a guard page too).
    stack: Option<KernelStack>,
    /// Level-4 page table to run with (the kernel's, or a process's own).
    cr3: PhysFrame,
    /// The user thread (and through it the process) this kernel thread
    /// runs (None = pure kernel thread).
    task: Option<Arc<Task>>,
    state: State,
    /// Some CPU is (still) running on this thread's stack.
    on_cpu: bool,
    /// A wake-up arrived while it was running: its next block returns at once.
    woken: bool,
    /// One of the per-CPU idle threads (never in the ready queue).
    idle: bool,
    /// CPUs that have flushed this thread's stack from their TLB (stack.rs).
    tlb_clean_on: u32,
    /// CPUs this thread has run on, one bit each (for tests and curiosity).
    ran_on: u32,
}

struct Scheduler {
    /// All threads. Boxed so each Thread has a fixed address while the map
    /// changes: `switch_context` gets a raw pointer to `rsp`.
    threads: BTreeMap<Tid, Box<Thread>>,
    /// Threads waiting for a CPU, in order.
    ready: VecDeque<Tid>,
    /// Per CPU: the running thread, the idle thread, and the thread we just
    /// switched away from (for `finish_switch`).
    current: [Option<Tid>; MAX_CPUS],
    idle: [Option<Tid>; MAX_CPUS],
    prev: [Option<Tid>; MAX_CPUS],
    next_id: Tid,
}

static SCHEDULER: Mutex<Scheduler> = Mutex::new(Scheduler {
    threads: BTreeMap::new(),
    ready: VecDeque::new(),
    current: [None; MAX_CPUS],
    idle: [None; MAX_CPUS],
    prev: [None; MAX_CPUS],
    next_id: 0,
});

/// Lock the scheduler. Interrupts must be off (locking rule above).
fn lock() -> spin::MutexGuard<'static, Scheduler> {
    debug_assert!(!interrupts::are_enabled(), "scheduler lock with interrupts on");
    SCHEDULER.lock()
}

/// Per-CPU timer statistics: ticks seen, and ticks that found the CPU busy.
static CPU_TICKS: [AtomicU64; MAX_CPUS] = [const { AtomicU64::new(0) }; MAX_CPUS];
static CPU_BUSY_TICKS: [AtomicU64; MAX_CPUS] = [const { AtomicU64::new(0) }; MAX_CPUS];
/// Is CPU n running its idle thread right now?
static CPU_IDLING: [AtomicBool; MAX_CPUS] = [const { AtomicBool::new(false) }; MAX_CPUS];
/// Kernel stack overflows caught so far.
static STACK_OVERFLOWS: AtomicU64 = AtomicU64::new(0);
/// Woken whenever a thread is gone for good (for `join`).
static THREAD_EXITED: WaitQueue = WaitQueue::new();

impl Scheduler {
    fn thread(&mut self, id: Tid) -> &mut Thread {
        self.threads.get_mut(&id).expect("no such thread")
    }

    fn add(&mut self, thread: Thread) -> Tid {
        let id = self.next_id;
        self.next_id += 1;
        self.threads.insert(id, Box::new(thread));
        id
    }

    /// The thread running on `cpu`. On first use on CPU 0 that's the boot
    /// code, which we adopt as thread "main" (its rsp is saved on its first
    /// switch). APs are registered in `run_ap` before they schedule.
    fn current(&mut self, cpu: usize) -> Tid {
        if let Some(id) = self.current[cpu] {
            return id;
        }
        let mut main = new_thread("main", None, Cr3::read().0, None);
        main.state = State::Running;
        main.on_cpu = true;
        let id = self.add(main);
        self.current[cpu] = Some(id);
        id
    }

    /// Next runnable thread from the queue, skipping threads another CPU is
    /// still switching away from (they stay queued).
    fn pop_ready(&mut self) -> Option<Tid> {
        for _ in 0..self.ready.len() {
            let id = self.ready.pop_front()?;
            match self.threads.get(&id) {
                Some(t) if t.state == State::Ready && !t.on_cpu => return Some(id),
                Some(t) if t.state == State::Ready => self.ready.push_back(id),
                _ => {} // stale entry (can't really happen): drop it
            }
        }
        None
    }

    /// Make a blocked/sleeping thread ready. Returns true if it was queued.
    fn make_ready(&mut self, id: Tid) -> bool {
        let Some(t) = self.threads.get_mut(&id) else { return false };
        match t.state {
            State::Blocked | State::Sleeping(_) => {
                t.state = State::Ready;
                self.ready.push_back(id);
                true
            }
            // Not asleep (yet): it may be just about to block. Make sure it
            // doesn't, or this wake-up would be lost.
            State::Running | State::Ready => {
                t.woken = true;
                false
            }
            State::Finished => false,
        }
    }
}

fn new_thread(name: &str, stack: Option<KernelStack>, cr3: PhysFrame, task: Option<Arc<Task>>) -> Thread {
    Thread {
        name: String::from(name),
        rsp: 0,
        stack,
        cr3,
        task,
        state: State::Ready,
        on_cpu: false,
        woken: false,
        idle: false,
        tlb_clean_on: 0,
        ran_on: 0,
    }
}

/// Create CPU 0's idle thread (smp::init calls this). Before this, a CPU
/// with nothing to run just halts inside whichever thread is blocking.
pub fn init() {
    let id = new_kernel_thread("idle0", idle_entry, crate::memory::kernel_l4_frame(), None, false);
    interrupts::without_interrupts(|| {
        let mut s = lock();
        s.current(0);
        s.thread(id).idle = true;
        s.idle[0] = Some(id);
    });
}

fn idle_entry() {
    idle_loop();
}

/// An AP (from smp.rs) turns its boot code into its idle thread, on the
/// boot stack it was given. Interrupts are off.
pub fn run_ap(cpu: usize, stack: KernelStack) -> ! {
    {
        let mut s = lock();
        let mut idle = new_thread("idle", Some(stack), crate::memory::kernel_l4_frame(), None);
        idle.name = alloc::format!("idle{}", cpu);
        idle.idle = true;
        idle.state = State::Running;
        idle.on_cpu = true;
        idle.tlb_clean_on = 1 << cpu;
        let id = s.add(idle);
        s.current[cpu] = Some(id);
        s.idle[cpu] = Some(id);
    }
    CPU_IDLING[cpu].store(true, Ordering::Relaxed);
    idle_loop()
}

/// What a CPU does when there's nothing to do: run whatever becomes ready,
/// otherwise sleep (`hlt`) until an interrupt.
fn idle_loop() -> ! {
    loop {
        interrupts::disable();
        schedule(); // returns when the queue is empty (again)
        let cpu = smp::cpu_id();
        // Say "I'm idle" *before* the last look at the queue: whoever queues
        // work after that look sees the flag and sends us an IPI.
        smp::set_idle(cpu, true);
        if interrupts::without_interrupts(|| lock().ready.is_empty()) {
            interrupts::enable_and_hlt(); // `sti; hlt` with no gap for a lost IPI
        } else {
            interrupts::enable();
        }
        smp::set_idle(cpu, false);
    }
}

/// Start a new kernel thread running `entry`. Needs the heap and memory.
pub fn spawn(name: &str, entry: fn()) -> Tid {
    spawn_with(name, entry, crate::memory::kernel_l4_frame(), None)
}

/// Start a kernel thread for user thread `task`: runs `entry` (in the
/// kernel) with the process's address space `cr3` active; `entry` then
/// drops to ring 3.
pub fn spawn_with(name: &str, entry: fn(), cr3: PhysFrame, task: Option<Arc<Task>>) -> Tid {
    new_kernel_thread(name, entry, cr3, task, true)
}

fn new_kernel_thread(name: &str, entry: fn(), cr3: PhysFrame, task: Option<Arc<Task>>, queue: bool) -> Tid {
    let stack = KernelStack::new().expect("out of memory for a kernel stack");

    // Build the initial stack, top down (stacks grow towards lower addresses):
    //   [return address = thread_trampoline]
    //   rbp, rbx, r12 (= entry), r13, r14, r15     <- rsp points at r15
    // exactly what `switch_context` pops before its `ret`.
    let top = stack.top();
    let initial: [u64; 7] = [
        0,                                // r15
        0,                                // r14
        0,                                // r13
        entry as *const () as u64,        // r12: picked up by the trampoline
        0,                                // rbx
        0,                                // rbp
        thread_trampoline as *const () as u64, // return address
    ];
    let rsp = top - 8 * initial.len() as u64;
    unsafe { core::ptr::copy_nonoverlapping(initial.as_ptr(), rsp as *mut u64, initial.len()) };

    let mut thread = new_thread(name, Some(stack), cr3, task);
    thread.rsp = rsp;
    let id = interrupts::without_interrupts(|| {
        let mut s = lock();
        s.current(smp::cpu_id()); // adopt "main" if this is the first spawn
        let id = s.add(thread);
        if queue {
            s.ready.push_back(id);
        }
        id
    });
    if queue {
        smp::kick_idle_cpu();
    }
    id
}

/// Switch this CPU to the next thread in the queue, if there is one.
/// Called with interrupts off: from the timer interrupt, `yield_now`, or a
/// thread that just blocked. Returns when this thread runs again (maybe on
/// another CPU!), or right away if there's nothing else to run.
pub fn schedule() {
    let cpu = smp::cpu_id();
    let (old_rsp, new_rsp) = {
        let mut s = lock();
        let Some(cur) = s.current[cpu] else { return }; // no threads yet
        let cur_state = s.thread(cur).state;
        let cur_idle = s.thread(cur).idle;
        let next = match s.pop_ready() {
            Some(next) => next,
            // Nobody waiting: a running thread just carries on...
            None if cur_state == State::Running => return,
            // ...one that blocked or finished hands over to the idle thread.
            None => match s.idle[cpu] {
                Some(idle) if idle != cur => idle,
                _ => return, // no idle thread yet: the caller halts instead
            },
        };
        // Back of the line for the current thread (if it can still run).
        if cur_state == State::Running {
            if cur_idle {
                s.thread(cur).state = State::Ready; // idle threads aren't queued
            } else {
                s.thread(cur).state = State::Ready;
                s.ready.push_back(cur);
            }
        }

        let bit = 1u32 << cpu;
        let t = s.thread(next);
        t.state = State::Running;
        t.on_cpu = true;
        t.ran_on |= bit;
        if let Some(task) = &t.task {
            task.process.ran_on.fetch_or(bit, core::sync::atomic::Ordering::Relaxed);
            // Thread-local storage: each user thread has its own FS base
            // (set with arch_prctl). The CPU has one FS_BASE register, so
            // it's reloaded on every switch to a user thread.
            x86_64::registers::model_specific::FsBase::write(x86_64::VirtAddr::new_truncate(
                task.fs_base.load(core::sync::atomic::Ordering::Relaxed),
            ));
        }
        if let Some(stack) = &t.stack {
            if t.tlb_clean_on & bit == 0 {
                stack.flush_local_tlb(); // see "Reusing a slot" in stack.rs
                t.tlb_clean_on |= bit;
            }
            // Where interrupts (TSS.RSP0) and `syscall` from ring 3 land.
            crate::gdt::set_kernel_stack(stack.top());
            crate::syscall::set_kernel_stack(stack.top());
        }
        if Cr3::read().0 != t.cr3 {
            // Switch address space. Safe for the code running right now:
            // all kernel mappings (incl. every kernel stack) are identical in
            // every address space. Writing CR3 also flushes the TLB.
            unsafe { Cr3::write(t.cr3, Cr3::read().1) };
        }
        // Which address space this CPU runs now (TLB shootdowns, memory.rs).
        smp::set_active_cr3(cpu, t.cr3.start_address().as_u64());
        CPU_IDLING[cpu].store(t.idle, Ordering::Relaxed);
        let new_rsp = t.rsp;
        s.current[cpu] = Some(next);
        s.prev[cpu] = Some(cur);
        let old_rsp: *mut u64 = &mut s.thread(cur).rsp;
        (old_rsp, new_rsp)
        // The lock is released here, *before* switching: the next thread
        // will never come back to release a lock we still held.
    };
    unsafe { switch_context(old_rsp, new_rsp) };
    finish_switch();
}

/// Right after a switch, on the new thread: the previous thread's stack is
/// no longer in use. Clear its `on_cpu` so other CPUs may run it, and if it
/// had finished, remove it and free its stack (and, if it was the last user
/// of its process, the process's memory, via `Drop for Process`).
fn finish_switch() {
    let cpu = smp::cpu_id();
    let (leftovers, requeued) = {
        let mut s = lock();
        let Some(prev) = s.prev[cpu].take() else { return };
        let t = s.thread(prev);
        t.on_cpu = false;
        let state = t.state;
        // A finished thread: take out what it owns (its stack, and its
        // share of the process) to free below, outside the lock.
        let leftovers = (state == State::Finished).then(|| (prev, t.stack.take(), t.task.take()));
        (leftovers, state == State::Ready)
    };
    if requeued {
        smp::kick_idle_cpu(); // it may have been skipped while on_cpu was set
    }
    if let Some((dead, stack, process)) = leftovers {
        // Free first, *then* remove it from the table: `join` waits for the
        // removal, so once `join` returns the memory really is free (the
        // leak tests count frames right after joining).
        drop(stack); // with interrupts still off: the allocators allow that
        drop(process);
        lock().threads.remove(&dead);
        THREAD_EXITED.wake_all();
    }
}

/// Every CPU's APIC timer tick: count it for the statistics, then preempt.
pub fn timer_tick() {
    let cpu = smp::cpu_id();
    CPU_TICKS[cpu].fetch_add(1, Ordering::Relaxed);
    if !CPU_IDLING[cpu].load(Ordering::Relaxed) {
        CPU_BUSY_TICKS[cpu].fetch_add(1, Ordering::Relaxed);
    }
    schedule();
}

/// (ticks, busy ticks) of CPU `cpu`'s APIC timer so far.
pub fn cpu_ticks(cpu: usize) -> (u64, u64) {
    (CPU_TICKS[cpu].load(Ordering::Relaxed), CPU_BUSY_TICKS[cpu].load(Ordering::Relaxed))
}

/// Give up the rest of this time slice (also callable outside interrupts).
pub fn yield_now() {
    interrupts::without_interrupts(schedule);
}

/// The process the running thread belongs to, if any.
pub fn current_process() -> Option<Arc<Process>> {
    interrupts::without_interrupts(|| {
        let mut s = lock();
        let cur = s.current(smp::cpu_id());
        s.thread(cur).task.as_ref().map(|t| t.process.clone())
    })
}

/// The user thread the running kernel thread carries, if any.
pub fn current_task() -> Option<Arc<Task>> {
    interrupts::without_interrupts(|| {
        let mut s = lock();
        let cur = s.current(smp::cpu_id());
        s.thread(cur).task.clone()
    })
}

/// End the running thread: mark it finished and switch away for good. The
/// next thread on this CPU frees its stack (see `finish_switch`).
pub fn exit_current() -> ! {
    interrupts::disable();
    {
        let mut s = lock();
        let cur = s.current(smp::cpu_id());
        s.thread(cur).state = State::Finished;
    }
    loop {
        schedule(); // never returns once there's an idle thread
        // No idle thread (scheduler not fully started): wait for anything.
        interrupts::enable_and_hlt();
        interrupts::disable();
    }
}

/// Name of the running thread (for printing).
pub fn current_name() -> String {
    interrupts::without_interrupts(|| {
        let mut s = lock();
        let cur = s.current(smp::cpu_id());
        s.thread(cur).name.clone()
    })
}

/// Id of the running thread.
pub fn current_id() -> Tid {
    interrupts::without_interrupts(|| lock().current(smp::cpu_id()))
}

/// Is the scheduler fully running (CPU 0 has its idle thread)?
pub fn scheduler_started() -> bool {
    interrupts::without_interrupts(|| lock().idle[0].is_some())
}

/// Does thread `id` still exist (not finished, or finished but its stack
/// not freed yet)?
pub fn exists(id: Tid) -> bool {
    interrupts::without_interrupts(|| lock().threads.contains_key(&id))
}

/// Wait until thread `id` is completely gone (finished and freed).
pub fn join(id: Tid) {
    THREAD_EXITED.wait_until(|| (!lock().threads.contains_key(&id)).then_some(()));
}

/// Number of threads (including idle threads and "main").
pub fn thread_count() -> usize {
    interrupts::without_interrupts(|| lock().threads.len())
}

/// CPUs thread `id` has run on so far (bit n = CPU n), if it exists.
pub fn ran_on(id: Tid) -> Option<u32> {
    interrupts::without_interrupts(|| lock().threads.get(&id).map(|t| t.ran_on))
}

/// How many kernel stack overflows have been caught.
pub fn stack_overflows() -> u64 {
    STACK_OVERFLOWS.load(Ordering::Relaxed)
}

/// Make thread `id` runnable again (if it's blocked or sleeping), and poke
/// an idle CPU to run it.
pub fn wake(id: Tid) {
    if interrupts::without_interrupts(|| lock().make_ready(id)) {
        smp::kick_idle_cpu();
    }
}

/// PIT tick (CPU 0's timer interrupt): wake every thread whose sleep is over.
pub fn wake_sleepers(now: u64) {
    let woke = {
        let mut s = lock();
        let due: Vec<Tid> = s
            .threads
            .iter()
            .filter(|(_, t)| matches!(t.state, State::Sleeping(at) if at <= now))
            .map(|(id, _)| *id)
            .collect();
        for &id in &due {
            s.make_ready(id);
        }
        !due.is_empty()
    };
    if woke {
        smp::kick_idle_cpu();
    }
}

/// Block the running thread (until `wake`, or until PIT tick `until`), and
/// run other threads meanwhile. Must be called with interrupts *off*, so a
/// wake-up can't slip in between the caller checking its condition and us
/// blocking; returns with interrupts still off. Returns at once if a
/// wake-up already arrived (`woken`).
fn block_current(until: Option<u64>) {
    let me = {
        let mut s = lock();
        let me = s.current(smp::cpu_id());
        let t = s.thread(me);
        if t.woken {
            t.woken = false;
            return;
        }
        t.state = match until {
            Some(at) => State::Sleeping(at),
            None => State::Blocked,
        };
        me
    };
    loop {
        schedule(); // returns once we've been woken and picked again
        if done_waiting(me) {
            return;
        }
        // We're still here: there was nothing (not even an idle thread) to
        // switch to. Halt until an interrupt possibly wakes us.
        // `enable_and_hlt` = `sti; hlt` with no gap for the interrupt.
        interrupts::enable_and_hlt();
        interrupts::disable();
        if done_waiting(me) {
            return;
        }
    }
}

/// After `schedule` returned to a thread that blocked: has it been woken?
/// If it was woken but never left the CPU, it's Ready and queued: take it
/// back out of the queue, it's running right here.
fn done_waiting(me: Tid) -> bool {
    let mut s = lock();
    match s.thread(me).state {
        State::Blocked | State::Sleeping(_) => false,
        State::Ready => {
            s.thread(me).state = State::Running;
            s.ready.retain(|&id| id != me);
            true
        }
        _ => true,
    }
}

/// Clear a stale wake-up flag before starting to wait for something new.
fn clear_woken() {
    let mut s = lock();
    let me = s.current(smp::cpu_id());
    s.thread(me).woken = false;
}

/// Sleep for `ticks` PIT ticks (~55 ms each). The thread is off the run
/// queue meanwhile, so it costs no CPU time; the PIT handler wakes it.
pub fn sleep_ticks(ticks: u64) {
    let was_enabled = interrupts::are_enabled();
    interrupts::disable();
    let wake_at = crate::interrupts::ticks() + ticks.max(1);
    clear_woken();
    while crate::interrupts::ticks() < wake_at {
        block_current(Some(wake_at)); // may return early on a stray wake-up
    }
    if was_enabled {
        interrupts::enable();
    }
}

/// `sleep_ticks` that a signal can cut short: `Err(ERESTARTNOHAND)` then
/// (nanosleep and pause are never restarted after a handler ran).
pub fn sleep_ticks_intr(ticks: u64) -> Result<(), Errno> {
    let was_enabled = interrupts::are_enabled();
    interrupts::disable();
    let wake_at = crate::interrupts::ticks() + ticks.max(1);
    clear_woken();
    let mut result = Ok(());
    while crate::interrupts::ticks() < wake_at {
        if crate::signal::interrupted() {
            result = Err(Errno::ERESTARTNOHAND);
            break;
        }
        block_current(Some(wake_at));
    }
    if was_enabled {
        interrupts::enable();
    }
    result
}

/// Block until a signal arrives (pause()).
pub fn pause() -> Errno {
    static NEVER: WaitQueue = WaitQueue::new();
    match NEVER.wait_until_intr(|| None::<()>) {
        Err(e) => if e == Errno::ERESTARTSYS { Errno::ERESTARTNOHAND } else { e },
        Ok(()) => unreachable!(),
    }
}

/// The running thread overflowed its kernel stack (the page-fault or
/// double-fault handler saw an access to a guard page). Say so, and kill
/// it: its process if it has one, else just the thread.
pub fn die_of_stack_overflow(addr: u64) -> ! {
    STACK_OVERFLOWS.fetch_add(1, Ordering::Relaxed);
    let name = current_name();
    crate::kprintln!(
        "[kernel] kernel stack overflow in thread '{}' (hit guard page {:#x}); killing it",
        name,
        addr & !0xfff
    );
    if current_process().is_some() {
        crate::process::exit_group(crate::process::ExitStatus::Killed(crate::signal::SIGSEGV, "kernel stack overflow"));
    }
    exit_current()
}

/// Threads waiting for some event (like a Linux `wait_queue_head_t`).
///
/// Usage: `queue.wait_until(|| condition)` blocks until the closure returns
/// `Some(value)`; whoever makes the condition true calls `queue.wake_all()`.
/// The closure runs with interrupts off, so it must be quick and only touch
/// locks that follow the locking rule at the top of this file.
pub struct WaitQueue {
    waiting: Mutex<Vec<Tid>>,
}

impl WaitQueue {
    pub const fn new() -> Self {
        WaitQueue { waiting: Mutex::new(Vec::new()) }
    }

    pub fn wait_until<T>(&self, condition: impl FnMut() -> Option<T>) -> T {
        match self.wait(None, false, condition) {
            Ok(Some(value)) => value,
            _ => unreachable!("no deadline, not interruptible"),
        }
    }

    /// Like `wait_until`, but give up at PIT tick `deadline`: returns
    /// `None` if the condition still wasn't true by then. (Network code uses
    /// this for timeouts: connect, DNS, ping, `SO_RCVTIMEO`.)
    pub fn wait_until_deadline<T>(&self, deadline: u64, condition: impl FnMut() -> Option<T>) -> Option<T> {
        self.wait(Some(deadline), false, condition).ok().flatten()
    }

    /// Like `wait_until`, but a signal for the process interrupts the wait:
    /// `Err(ERESTARTSYS)`. Every wait a program can get stuck in (reading
    /// the keyboard or a pipe, wait4, sockets, ...) uses this, or Ctrl+C
    /// couldn't stop a program blocked in `read`. On the way back to user
    /// space the signal code turns ERESTARTSYS into EINTR or restarts the
    /// call (see signal.rs).
    pub fn wait_until_intr<T>(&self, condition: impl FnMut() -> Option<T>) -> Result<T, Errno> {
        self.wait(None, true, condition).map(|v| v.expect("no deadline"))
    }

    /// Interruptible *and* with a deadline (`Ok(None)` = timed out).
    pub fn wait_until_deadline_intr<T>(&self, deadline: u64, condition: impl FnMut() -> Option<T>) -> Result<Option<T>, Errno> {
        self.wait(Some(deadline), true, condition)
    }

    fn wait<T>(&self, deadline: Option<u64>, interruptible: bool, mut condition: impl FnMut() -> Option<T>) -> Result<Option<T>, Errno> {
        let was_enabled = interrupts::are_enabled();
        let me = current_id();
        let result = loop {
            interrupts::disable();
            // The order matters (the classic "lost wake-up" bug): put
            // ourselves on the list *first*, then check. If the condition
            // becomes true on another CPU after our check, its wake_all()
            // finds us on the list and sets `woken`, so block_current
            // returns straight away instead of sleeping forever. (Signals
            // wake us with `wake(tid)` directly, which works the same way.)
            clear_woken();
            self.waiting.lock().push(me);
            if let Some(value) = condition() {
                break Ok(Some(value));
            }
            if interruptible && crate::signal::interrupted() {
                break Err(Errno::ERESTARTSYS);
            }
            if deadline.is_some_and(|d| crate::interrupts::ticks() >= d) {
                break Ok(None);
            }
            // Sleep until woken (or the deadline tick, whichever is first).
            block_current(deadline);
            // Woken up: take us off the list (a timer or signal wake-up
            // leaves us on it), then check again -- somebody else may have
            // consumed what we were waiting for.
            self.waiting.lock().retain(|&id| id != me);
        };
        self.waiting.lock().retain(|&id| id != me);
        if was_enabled {
            interrupts::enable();
        }
        result
    }

    pub fn wake_all(&self) {
        let waiting: Vec<Tid> = interrupts::without_interrupts(|| self.waiting.lock().drain(..).collect());
        for id in waiting {
            wake(id);
        }
    }
}

// --- assembly ---------------------------------------------------------------

unsafe extern "C" {
    /// Save callee-saved registers + rsp into `*old_rsp`, load `new_rsp`,
    /// restore that thread's registers and return into it.
    fn switch_context(old_rsp: *mut u64, new_rsp: u64);
    /// First "return address" of every new thread.
    fn thread_trampoline();
}

// System V calling convention: 1st argument in rdi, 2nd in rsi. A normal
// function call already saved the caller-saved registers, so only the
// callee-saved ones (rbx, rbp, r12-r15) need saving here.
global_asm!(
    ".global switch_context",
    "switch_context:",
    "    push rbp",
    "    push rbx",
    "    push r12",
    "    push r13",
    "    push r14",
    "    push r15",
    "    mov [rdi], rsp", // old thread's rsp -> *old_rsp
    "    mov rsp, rsi",   // switch stacks!
    "    pop r15",
    "    pop r14",
    "    pop r13",
    "    pop r12",
    "    pop rbx",
    "    pop rbp",
    "    ret",            // pops the new thread's return address
    "",
    ".global thread_trampoline",
    "thread_trampoline:",
    "    mov rdi, r12",   // entry function pointer -> 1st argument
    "    and rsp, -16",   // ABI: stack 16-byte aligned before `call`
    "    call {start}",
    "    ud2",            // thread_start never returns
    start = sym thread_start,
);

/// Rust side of a new thread's life.
extern "C" fn thread_start(entry: usize) -> ! {
    // We arrived here from inside `schedule`, like any switched-to thread:
    // do the same clean-up as `schedule` does after its switch.
    finish_switch();
    // Interrupts are off in `schedule`. Turn them back on, or this thread
    // could never be preempted.
    interrupts::enable();

    let entry: fn() = unsafe { core::mem::transmute(entry) };
    entry();

    // The thread function returned: mark it finished so the scheduler skips it.
    exit_current();
}