361 lines
13 KiB
Plaintext
361 lines
13 KiB
Plaintext
import drivers.serial.uart.serial as serial
|
||
import sched.process as proc
|
||
import mm.mm as mm
|
||
import intr.gdt as gdt
|
||
import drivers.core.cpu.cpu as cpu
|
||
import viperlib
|
||
import asm
|
||
|
||
define THREAD_READY = 0
|
||
define THREAD_RUNNING = 1
|
||
define THREAD_BLOCKED = 2
|
||
define THREAD_DONE = 3
|
||
define MAX_THREADS = 16
|
||
define THREAD_STACK_SIZE = 65536
|
||
define SCHED_QUANTUM = 10
|
||
|
||
_yield_cnt: uint32 = 0
|
||
_yield_lock: int = 0
|
||
|
||
|
||
class Thread:
|
||
rsp: uint64
|
||
stack: rawptr
|
||
stack_size: uint64
|
||
func: rawptr
|
||
arg: rawptr
|
||
state: int
|
||
tid: int
|
||
pid: int
|
||
# Ring3 → Ring0 的落栈点（TSS.RSP0 / gs:[8]）：用户线程才非 0，内核线程为 0。
|
||
# 必须按线程切换 —— 同进程多用户线程共用进程级内核栈会互相压栈帧。
|
||
kern_stack: uint64
|
||
|
||
def block(self):
|
||
self.state = THREAD_BLOCKED
|
||
Scheduler._yield()
|
||
|
||
def unblock(self):
|
||
if self.state == THREAD_BLOCKED:
|
||
self.state = THREAD_READY
|
||
|
||
@property
|
||
def is_running(self) -> bool: return self.state == THREAD_RUNNING
|
||
@property
|
||
def is_blocked(self) -> bool: return self.state == THREAD_BLOCKED
|
||
@property
|
||
def is_done(self) -> bool: return self.state == THREAD_DONE
|
||
|
||
|
||
class Scheduler:
|
||
threads: [MAX_THREADS]Thread
|
||
current_tid: int
|
||
thread_count: int
|
||
enabled: int
|
||
needs_reschedule: int
|
||
tick_count: int
|
||
|
||
def __init__(self):
|
||
global _sched_ptr
|
||
_sched_ptr = &_sched_storage
|
||
s: Scheduler* = _sched_ptr
|
||
s.current_tid = 0
|
||
s.thread_count = 1
|
||
s.enabled = 0
|
||
s.needs_reschedule = 0
|
||
s.tick_count = 0
|
||
for i in range(MAX_THREADS):
|
||
s.threads[i].rsp = 0
|
||
s.threads[i].stack = None
|
||
s.threads[i].stack_size = 0
|
||
s.threads[i].func = None
|
||
s.threads[i].arg = None
|
||
s.threads[i].state = THREAD_DONE
|
||
s.threads[i].tid = i
|
||
s.threads[i].pid = 0
|
||
s.enabled = 1
|
||
s.threads[0].state = THREAD_RUNNING
|
||
s.threads[0].tid = 0
|
||
main_rsp_ptr: rawptr = &s.threads[0].rsp
|
||
__asm__(f"""mov qword ptr [{asminp(main_rsp_ptr, ASM_DESCR.REG_ANY)}], rsp""",
|
||
op=[ASM_DESCR.CLOBBER_MEMORY])
|
||
# fail-fast：_thread_entry 的 asm 以硬编码偏移访问线程表（sizeof(Thread)=64、
|
||
# current_tid=+1024=16*64）。增删 Thread 字段会让它们失配 ⇒ 线程入口会按错误的
|
||
# 偏移取 func/arg 并 call 任意地址。这里直接核对，不符立即停机暴露。
|
||
thr0_addr: uint64 = uint64(&s.threads[0])
|
||
curt_addr: uint64 = uint64(&s.current_tid)
|
||
if Thread.__sizeof__() != 64 or (curt_addr - thr0_addr) != 1024:
|
||
serial.puts("[sched] FATAL: Thread 布局与 _thread_entry 硬编码偏移(64/1024)不一致\n")
|
||
while True:
|
||
asm.sti()
|
||
asm.hlt()
|
||
|
||
@staticmethod
|
||
def _yield():
|
||
global _switch_old_rsp, _switch_new_rsp, _yield_cnt, _yield_lock
|
||
_if_saved: int = 0
|
||
__asm__(f"""pushfq
|
||
pop rax
|
||
shr rax, 9
|
||
and eax, 1
|
||
mov {asmout(_if_saved, ASM_DESCR.OUTPUT_REG)}, eax""",
|
||
op=[ASM_DESCR.CLOBBER_RAX])
|
||
asm.cli()
|
||
if _yield_lock:
|
||
if _if_saved:
|
||
asm.sti()
|
||
return
|
||
_yield_lock = 1
|
||
s: Scheduler* = _sched_ptr
|
||
s.needs_reschedule = 0
|
||
s.tick_count = 0
|
||
old_tid: int = s.current_tid
|
||
next_tid: int = old_tid + 1
|
||
if next_tid >= s.thread_count:
|
||
next_tid = 0
|
||
found: int = 0
|
||
for count in range(s.thread_count):
|
||
st: int = s.threads[next_tid].state
|
||
if st == THREAD_READY or st == THREAD_RUNNING:
|
||
found = 1
|
||
break
|
||
next_tid += 1
|
||
if next_tid >= s.thread_count:
|
||
next_tid = 0
|
||
if found == 0:
|
||
_yield_lock = 0
|
||
if _if_saved:
|
||
asm.sti()
|
||
return
|
||
if next_tid == old_tid:
|
||
_yield_lock = 0
|
||
if _if_saved:
|
||
asm.sti()
|
||
return
|
||
if s.threads[old_tid].state == THREAD_RUNNING:
|
||
s.threads[old_tid].state = THREAD_READY
|
||
s.threads[next_tid].state = THREAD_RUNNING
|
||
_switch_old_rsp = &s.threads[old_tid].rsp
|
||
_switch_new_rsp = &s.threads[next_tid].rsp
|
||
s.current_tid = next_tid
|
||
next_pid: int = s.threads[next_tid].pid
|
||
_yield_cnt += 1
|
||
proc.switch_pid(next_pid)
|
||
# 按**线程**设置 Ring3→Ring0 的落栈点：switch_pid 只按进程设置（会把全进程的用户线程
|
||
# 指到同一个内核栈上），两个线程都在 syscall 里 yield 时后者的栈帧会压掉前者的。
|
||
next_ks: uint64 = s.threads[next_tid].kern_stack
|
||
if next_ks != 0:
|
||
gdt.set_tss_rsp0(next_ks)
|
||
cpu.set_kernel_rsp0(next_ks)
|
||
_yield_lock = 0
|
||
_do_switch()
|
||
if _if_saved:
|
||
asm.sti()
|
||
|
||
@staticmethod
|
||
def sched_tick():
|
||
s: Scheduler* = _sched_ptr
|
||
if not s.enabled: return
|
||
if s.thread_count <= 1: return
|
||
s.tick_count += 1
|
||
if s.tick_count >= SCHED_QUANTUM:
|
||
s.tick_count = 0
|
||
s.needs_reschedule = 1
|
||
|
||
@staticmethod
|
||
def try_reschedule():
|
||
s: Scheduler* = _sched_ptr
|
||
if s.needs_reschedule:
|
||
s.needs_reschedule = 0
|
||
Scheduler._yield()
|
||
|
||
@staticmethod
|
||
def disable():
|
||
s: Scheduler* = _sched_ptr
|
||
s.enabled = 0
|
||
|
||
@staticmethod
|
||
def enable():
|
||
s: Scheduler* = _sched_ptr
|
||
s.enabled = 1
|
||
|
||
@staticmethod
|
||
def current() -> Thread*:
|
||
s: Scheduler* = _sched_ptr
|
||
return &s.threads[s.current_tid]
|
||
|
||
@staticmethod
|
||
def get_thread(tid: int) -> Thread*:
|
||
s: Scheduler* = _sched_ptr
|
||
if tid < 0 or tid >= s.thread_count: return None
|
||
return &s.threads[tid]
|
||
|
||
@staticmethod
|
||
def create_thread(func: rawptr, arg: rawptr, thread_pid: int = 0) -> Thread*:
|
||
s: Scheduler* = _sched_ptr
|
||
if s.thread_count >= MAX_THREADS: return None
|
||
tid: int = s.thread_count
|
||
s.thread_count += 1
|
||
th: Thread* = &s.threads[tid]
|
||
th.state = THREAD_READY
|
||
th.tid = tid
|
||
th.pid = thread_pid
|
||
th.func = func
|
||
th.arg = arg
|
||
th.kern_stack = 0
|
||
stack_size: uint64 = THREAD_STACK_SIZE
|
||
stack_ptr: rawptr = mm.malloc(stack_size)
|
||
if not stack_ptr:
|
||
s.thread_count -= 1
|
||
return None
|
||
th.stack = stack_ptr
|
||
th.stack_size = stack_size
|
||
stack_top: uint64 = uint64(stack_ptr) + stack_size - 16
|
||
stack_top = stack_top & ~uint64(0xF)
|
||
entry_addr: uint64 = uint64(&_thread_entry)
|
||
sp: rawptr = rawptr(stack_top)
|
||
__asm__(f"""mov rax, {asminp(sp, ASM_DESCR.REG_ANY)}
|
||
mov qword ptr [rax - 64], 0x202
|
||
mov qword ptr [rax - 56], 0
|
||
mov qword ptr [rax - 48], 0
|
||
mov qword ptr [rax - 40], 0
|
||
mov qword ptr [rax - 32], 0
|
||
mov qword ptr [rax - 24], 0
|
||
mov qword ptr [rax - 16], 0
|
||
mov qword ptr [rax - 8], 0
|
||
mov qword ptr [rax], {asminp(entry_addr, ASM_DESCR.REG_ANY)}""",
|
||
op=[ASM_DESCR.CLOBBER_MEMORY, ASM_DESCR.CLOBBER_RAX])
|
||
stack_top = stack_top - 64
|
||
th.rsp = stack_top
|
||
return th
|
||
|
||
@staticmethod
|
||
def create_user_thread(thread_pid: int, entry: rawptr, arg: rawptr) -> Thread*:
|
||
"""创建**用户态**线程：入口在 Ring 3 执行（THREAD_CREATE 用）。
|
||
|
||
app 传来的函数指针是用户地址，绝不能用 create_thread 直接 call —— 那是在 Ring 0
|
||
执行用户代码，且与同进程其它线程共用一套内核栈。这里给线程配独立用户栈 + 独立
|
||
内核栈，由 _user_thread_entry 走 drop_to_user_mode 进 Ring 3。
|
||
"""
|
||
slot: int = proc.alloc_user_thread(thread_pid, entry, arg)
|
||
if slot < 0: return None
|
||
ctx: proc.UserThreadCtx* = proc.user_thread_ctx(slot)
|
||
if ctx is None: return None
|
||
th: Thread* = Scheduler.create_thread(&_user_thread_entry, rawptr(uint64(slot)), thread_pid)
|
||
if th is None: return None
|
||
th.kern_stack = ctx.kern_stack
|
||
ctx.tid = th.tid
|
||
return th
|
||
|
||
@staticmethod
|
||
def exit_current_thread():
|
||
"""用户线程显式退出（THREAD_EXIT）：回收资源 → 标记 DONE → 让出，永不返回。"""
|
||
s: Scheduler* = _sched_ptr
|
||
tid: int = s.current_tid
|
||
proc.free_user_thread(proc.find_user_slot_by_tid(tid))
|
||
if tid > 0 and tid < MAX_THREADS:
|
||
s.threads[tid].state = THREAD_DONE
|
||
Scheduler._yield()
|
||
# _yield 可能因"只有自己可运行"而直接返回：此后绝不能再跑用户代码（栈已回收）。
|
||
while True:
|
||
asm.sti()
|
||
asm.hlt()
|
||
Scheduler._yield()
|
||
|
||
|
||
_sched_storage: Scheduler = Scheduler()
|
||
export _sched_ptr: Scheduler* = &_sched_storage
|
||
export _switch_old_rsp: rawptr
|
||
export _switch_new_rsp: rawptr
|
||
|
||
@__attribute__(naked)
|
||
def _do_switch():
|
||
__asm__(f"""push rax
|
||
push rbx
|
||
push rbp
|
||
push r12
|
||
push r13
|
||
push r14
|
||
push r15
|
||
pushfq
|
||
mov rax, qword ptr [rip + _switch_old_rsp]
|
||
mov qword ptr [rax], rsp
|
||
mov rax, qword ptr [rip + _switch_new_rsp]
|
||
mov rsp, qword ptr [rax]
|
||
popfq
|
||
pop r15
|
||
pop r14
|
||
pop r13
|
||
pop r12
|
||
pop rbp
|
||
pop rbx
|
||
pop rax
|
||
ret""",
|
||
op=[ASM_DESCR.CLOBBER_MEMORY])
|
||
|
||
# Thread 布局：rsp0/stack8/stack_size16/func24/arg32/state40/tid44/pid48/kern_stack56
|
||
# ⇒ sizeof(Thread) = 64，Scheduler.threads[16] 占 1024 字节 ⇒ current_tid 在 +1024。
|
||
# 增删 Thread 字段时必须同步下面 asm 里的 1024 / 64（以及 +24/+32/+40 三个字段偏移）。
|
||
@__attribute__(naked)
|
||
def _thread_entry():
|
||
__asm__(f"""sub rsp, 8
|
||
mov rax, qword ptr [rip + _sched_ptr]
|
||
mov ecx, dword ptr [rax + 1024]
|
||
movsxd rcx, ecx
|
||
imul rcx, rcx, 64
|
||
mov rdi, qword ptr [rax + rcx + 32]
|
||
mov rax, qword ptr [rax + rcx + 24]
|
||
call rax
|
||
mov rax, qword ptr [rip + _sched_ptr]
|
||
mov ecx, dword ptr [rax + 1024]
|
||
movsxd rcx, ecx
|
||
imul rcx, rcx, 64
|
||
mov dword ptr [rax + rcx + 40], 3
|
||
add rsp, 8""",
|
||
# `call rax` 破坏全部 caller-saved 寄存器（x64 SysV）：除 RAX/RCX/RDX/
|
||
# R8-R11 外，RSI 与 RDI 也必须声明，否则编译器可能假设其值存活（N9）。
|
||
# RSP 由 `sub rsp,8` / `add rsp,8` 配对抵消，净变化为 0，无需声明。
|
||
op=[ASM_DESCR.CLOBBER_MEMORY, ASM_DESCR.CLOBBER_RAX,
|
||
ASM_DESCR.CLOBBER_RCX, ASM_DESCR.CLOBBER_RDX,
|
||
ASM_DESCR.CLOBBER_RSI, ASM_DESCR.CLOBBER_RDI,
|
||
ASM_DESCR.CLOBBER_R8, ASM_DESCR.CLOBBER_R9,
|
||
ASM_DESCR.CLOBBER_R10, ASM_DESCR.CLOBBER_R11])
|
||
Scheduler._yield()
|
||
while True:
|
||
asm.sti()
|
||
asm.hlt()
|
||
Scheduler._yield()
|
||
|
||
def _user_thread_entry(arg: rawptr) -> int:
|
||
"""用户态线程的内核侧引导：配好 TSS/GS 落栈点后 iretq 进 Ring 3，永不返回。"""
|
||
slot: int = int(uint64(arg))
|
||
ctx: proc.UserThreadCtx* = proc.user_thread_ctx(slot)
|
||
if ctx is None: return 0
|
||
# 与 execrunner/elf.vp 的 _elf_entry_wrapper 同款：直到 iretq 之前必须关中断
|
||
# （@pml4/TSS/GS 都是全局量，被定时器抢占时别的线程的 switch_pid 会把它们改掉）；
|
||
# iretq 用用户态帧里的 RFLAGS.IF(0x202) 重新开中断。
|
||
asm.cli()
|
||
proc.switch_pid(ctx.pid)
|
||
gdt.set_tss_rsp0(ctx.kern_stack)
|
||
cpu.set_kernel_rsp0(ctx.kern_stack)
|
||
proc.drop_to_user_mode(ctx.entry, ctx.user_stack_entry_rsp, ctx.kern_stack)
|
||
return 0
|
||
|
||
def thread_current() -> int:
|
||
s: Scheduler* = _sched_ptr
|
||
return s.current_tid
|
||
|
||
def needs_reschedule() -> int:
|
||
s: Scheduler* = _sched_ptr
|
||
return s.needs_reschedule
|
||
|
||
def sched_dump():
|
||
s: Scheduler* = _sched_ptr
|
||
dl: [80]char
|
||
viperlib.snprintf(&dl, 80, "[sched] dump: cur=%d cnt=%d yields=%u\n", s.current_tid, s.thread_count, _yield_cnt)
|
||
serial.puts(dl)
|
||
for i in range(s.thread_count):
|
||
sl: [80]char
|
||
viperlib.snprintf(&sl, 80, " t%d: st=%d pid=%d rsp=0x%lx\n", i, s.threads[i].state, s.threads[i].pid, s.threads[i].rsp)
|
||
serial.puts(sl)
|