Initial project import

tools/v/vlib/goroutines/scheduler.v

@@ -0,0 +1,642 @@
+// Copyright (c) 2019-2024 Alexander Medvednikov. All rights reserved.
+// Use of this source code is governed by an MIT license
+// that can be found in the LICENSE file.
+//
+// GMP Scheduler - translated from Go's runtime/proc.go.
+//
+// The scheduler distributes ready-to-run goroutines over worker threads.
+// Key functions translated from Go:
+//   - schedule()      -> scheduler_loop()
+//   - findRunnable()  -> find_runnable()
+//   - execute()       -> execute()
+//   - newproc()       -> goroutine_create()
+//   - wakep()         -> wake_p()
+//   - runqput()       -> runq_put()
+//   - runqget()       -> runq_get()
+//   - stealWork()     -> steal_work()
+module goroutines
+
+// import sync as _
+// import runtime as _
+
+// goroutine_create creates a new goroutine to run `f` with argument `arg`.
+// This is the equivalent of Go's `newproc` function.
+// Called by the compiler for `go expr()`.
+pub fn goroutine_create(f voidptr, arg voidptr, arg_size int) {
+	gp := newproc1(f, arg, arg_size)
+	// Put on global run queue so any M can pick it up.
+	// This is simpler than local queue + work stealing and avoids
+	// visibility issues when the creator is the main (non-scheduler) thread.
+	glob_runq_put(gp)
+	// Try to wake an idle P if needed
+	wake_p()
+}
+
+// newproc1 allocates and initializes a new G.
+// Translated from Go's newproc1 in proc.go.
+fn newproc1(f voidptr, arg voidptr, arg_size int) &Goroutine {
+	// Try to get a dead G from the local P's free list first
+	mut pp := get_current_p()
+	mut gp := if pp != unsafe { nil } {
+		gfget(mut pp)
+	} else {
+		unsafe { &Goroutine(nil) }
+	}
+	if gp == unsafe { nil } {
+		// Try global free list
+		gp = gfget_global()
+	}
+	if gp == unsafe { nil } {
+		// Allocate a new G
+		gp = &Goroutine{}
+	}
+
+	// Allocate or reuse stack
+	if gp.stack == unsafe { nil } {
+		stack_size := default_stack_size
+		gp.stack = unsafe { malloc(stack_size) }
+		gp.stack_size = stack_size
+	}
+
+	gp.fn_ptr = f
+	gp.fn_arg = arg
+	gp.status = .runnable
+	gp.preempt = false
+
+	// Assign goroutine ID
+	gp.id = assign_goid()
+
+	// Initialize context to run goroutine_entry
+	context_init(mut &gp.context, gp.stack, gp.stack_size, goroutine_trampoline, voidptr(gp))
+
+	// Track all goroutines
+	allgs_mu.acquire()
+	allgs << gp
+	allgs_mu.release()
+
+	return gp
+}
+
+// goroutine_trampoline is the entry point for all goroutines.
+// It calls the user function, then cleans up and returns to the scheduler.
+// Equivalent to Go's goexit -> goexit0.
+fn goroutine_trampoline(arg voidptr) {
+	mut gp := unsafe { &Goroutine(arg) }
+
+	// Call the actual function
+	if gp.fn_ptr != unsafe { nil } {
+		call_goroutine_fn(gp.fn_ptr, gp.fn_arg)
+	}
+
+	// Function returned - goroutine is done
+	goexit0(mut gp)
+}
+
+// call_goroutine_fn calls the goroutine's function pointer with the given argument.
+fn call_goroutine_fn(fn_ptr voidptr, arg voidptr) {
+	unsafe {
+		cb := GoFn(fn_ptr)
+		cb(arg)
+	}
+}
+
+// goexit0 handles goroutine cleanup after the user function returns.
+// Translated from Go's goexit0 in proc.go.
+fn goexit0(mut gp Goroutine) {
+	mut mp := get_current_m()
+
+	gp.status = .dead
+	gp.m = unsafe { nil }
+	gp.fn_ptr = unsafe { nil }
+	gp.fn_arg = unsafe { nil }
+	gp.wait_reason = ''
+	gp.preempt = false
+
+	// Dissociate G from M
+	if mp != unsafe { nil } {
+		mp.curg = unsafe { nil }
+	}
+
+	// Put the dead G on the free list for reuse
+	mut pp := get_current_p()
+	if pp != unsafe { nil } {
+		gfput(mut pp, gp)
+	} else {
+		gfput_global(gp)
+	}
+
+	// Switch back to the scheduler (g0 context).
+	// This returns to schedule_loop via execute's context_switch.
+	if mp != unsafe { nil } && mp.g0 != unsafe { nil } {
+		context_set(&mp.g0.context)
+	}
+}
+
+// schedule is the main scheduler entry point.
+// Finds a runnable goroutine and executes it. Never returns.
+// Translated from Go's schedule() in proc.go.
+pub fn schedule() {
+	mut mp := get_current_m()
+	if mp == unsafe { nil } {
+		return
+	}
+	mut pp := mp.p
+	if pp == unsafe { nil } {
+		return
+	}
+	pp.sched_tick++
+
+	// Find a runnable goroutine
+	gp, _ := find_runnable(mut mp, mut pp)
+	if gp == unsafe { nil } {
+		// No work found - park this M
+		park_m(mut mp)
+		return
+	}
+
+	// Execute the goroutine
+	execute(mut mp, mut pp, gp)
+}
+
+// find_runnable finds a runnable goroutine to execute.
+// Tries: local queue, global queue, work stealing from other P's.
+// Translated from Go's findRunnable() in proc.go.
+fn find_runnable(mut mp Machine, mut pp Processor) (&Goroutine, bool) {
+	// Check global queue every Nth tick for fairness (Go uses 61)
+	if pp.sched_tick % global_queue_check_interval == 0 {
+		gp := glob_runq_get()
+		if gp != unsafe { nil } {
+			return gp, false
+		}
+	}
+
+	// Try local run queue first
+	gp, inherit := runq_get(mut pp)
+	if gp != unsafe { nil } {
+		return gp, inherit
+	}
+
+	// Try global run queue
+	gp2 := glob_runq_get()
+	if gp2 != unsafe { nil } {
+		return gp2, false
+	}
+
+	// Try to steal work from other P's
+	gp3 := steal_work(mut pp)
+	if gp3 != unsafe { nil } {
+		return gp3, false
+	}
+
+	return unsafe { nil }, false
+}
+
+// execute runs a goroutine on the current M.
+// Translated from Go's execute() in proc.go.
+fn execute(mut mp Machine, mut pp Processor, gp &Goroutine) {
+	mut g := unsafe { gp }
+	// Associate G with M
+	mp.curg = g
+	g.m = mp
+	g.status = .running
+
+	// Switch context to the goroutine
+	context_switch(mut &mp.g0.context, &g.context)
+	// When we return here, the goroutine has yielded back to us
+	// The scheduler loop will be re-entered
+}
+
+// wake_p tries to wake an idle P to run goroutines.
+// Translated from Go's wakep() in proc.go.
+fn wake_p() {
+	// Check if there's an idle P
+	gsched.mu.acquire()
+	if gsched.npidle == 0 {
+		gsched.mu.release()
+		return
+	}
+	// Don't wake if there are already spinning M's
+	if gsched.nmspinning > 0 {
+		gsched.mu.release()
+		return
+	}
+	// Get an idle P
+	pp := pid_get()
+	if pp == unsafe { nil } {
+		gsched.mu.release()
+		return
+	}
+	gsched.mu.release()
+
+	// Start a new M for this P (or wake an idle one)
+	start_m(pp)
+}
+
+// park_m parks the current M - it goes to sleep waiting for work.
+fn park_m(mut mp Machine) {
+	// Release P and add M to idle list under lock
+	gsched.mu.acquire()
+	if mp.p != unsafe { nil } {
+		pid_put(mp.p)
+		mp.p = unsafe { nil }
+	}
+	mp.sched_link = gsched.midle
+	gsched.midle = mp
+	gsched.nmidle++
+	gsched.mu.release()
+
+	// Sleep until woken
+	mp.park.wait()
+
+	// Woken up - acquire a P and return to schedule_loop
+	acquire_p(mut mp)
+}
+
+// acquire_p tries to get an idle P for the given M.
+fn acquire_p(mut mp Machine) {
+	gsched.mu.acquire()
+	mut pp := pid_get()
+	gsched.mu.release()
+	if pp != unsafe { nil } {
+		wire_p(mut mp, mut pp)
+	}
+}
+
+// wire_p associates a P with an M.
+fn wire_p(mut mp Machine, mut pp Processor) {
+	mp.p = pp
+	pp.m = mp
+	pp.status = .running
+}
+
+// start_m starts or wakes an M to run the given P.
+// Translated from Go's startm() in proc.go.
+fn start_m(pp &Processor) {
+	gsched.mu.acquire()
+
+	// Try to get an idle M first
+	mut mp := gsched.midle
+	if mp != unsafe { nil } {
+		gsched.midle = mp.sched_link
+		gsched.nmidle--
+		gsched.mu.release()
+		// Give it the P and wake it
+		mut p := unsafe { pp }
+		wire_p(mut mp, mut p)
+		mp.park.post()
+		return
+	}
+
+	// No idle M available - create a new one
+	id := gsched.mnext
+	gsched.mnext++
+	gsched.mu.release()
+
+	new_m(id, pp)
+}
+
+// new_m creates a new OS thread (M) and associates it with a P.
+fn new_m(id i64, pp &Processor) {
+	mut mp := &Machine{
+		id: id
+		g0: &Goroutine{
+			status: .running
+		}
+	}
+	mut p := unsafe { pp }
+	wire_p(mut mp, mut p)
+	mp.thread = spawn m_thread_entry(mut mp)
+}
+
+// m_thread_entry is the entry point for new M (OS thread) goroutine scheduling loops.
+fn m_thread_entry(mut mp Machine) {
+	// Register this M as the current thread's Machine so that
+	// get_current_m()/get_current_p() work on worker threads.
+	set_current_m(mp)
+	// Enter the scheduling loop - this never returns
+	schedule_loop(mut mp)
+}
+
+// schedule_loop is the main loop for an M. It repeatedly finds and runs goroutines.
+fn schedule_loop(mut mp Machine) {
+	for {
+		if gsched.stopped {
+			return
+		}
+		mut pp := mp.p
+		if pp == unsafe { nil } {
+			acquire_p(mut mp)
+			pp = mp.p
+			if pp == unsafe { nil } {
+				// No P available, park
+				park_m(mut mp)
+				continue
+			}
+		}
+		pp.sched_tick++
+
+		gp, _ := find_runnable(mut mp, mut pp)
+		if gp == unsafe { nil } {
+			// No work - try spinning briefly before parking
+			if !mp.spinning {
+				mp.spinning = true
+				C.goroutines_atomic_fetch_add_i32(&gsched.nmspinning, 1)
+			}
+			// Spin a bit
+			mut found := false
+			for _ in 0 .. 20 {
+				gp2, _ := find_runnable(mut mp, mut pp)
+				if gp2 != unsafe { nil } {
+					mp.spinning = false
+					C.goroutines_atomic_fetch_sub_i32(&gsched.nmspinning, 1)
+					execute(mut mp, mut pp, gp2)
+					found = true
+					break
+				}
+				// Yield to avoid burning CPU
+				proc_yield(10)
+			}
+			if !found {
+				if mp.spinning {
+					mp.spinning = false
+					C.goroutines_atomic_fetch_sub_i32(&gsched.nmspinning, 1)
+				}
+				park_m(mut mp)
+			}
+			continue
+		}
+
+		if mp.spinning {
+			mp.spinning = false
+			C.goroutines_atomic_fetch_sub_i32(&gsched.nmspinning, 1)
+		}
+
+		execute(mut mp, mut pp, gp)
+	}
+}
+
+// steal_work tries to steal goroutines from other P's run queues.
+// Translated from Go's stealWork() in proc.go.
+fn steal_work(mut thisp Processor) &Goroutine {
+	n := gsched.allp.len
+	if n <= 1 {
+		return unsafe { nil }
+	}
+	// Randomize starting point to avoid contention
+	start := u32(C.rand()) % u32(n)
+	for i := u32(0); i < u32(n); i++ {
+		idx := (start + i) % u32(n)
+		pp := gsched.allp[idx]
+		if pp == thisp {
+			continue
+		}
+		// Try to steal half of the target's run queue
+		mut target := unsafe { pp }
+		gp := runq_steal(mut target, mut thisp)
+		if gp != unsafe { nil } {
+			return gp
+		}
+	}
+	return unsafe { nil }
+}
+
+// runq_steal steals half of pp's local run queue.
+// Translated from Go's runqgrab/runqsteal in proc.go.
+fn runq_steal(mut pp Processor, mut thisp Processor) &Goroutine {
+	t := C.goroutines_atomic_load_u32(&pp.runq_tail)
+	h := C.goroutines_atomic_load_u32(&pp.runq_head)
+	n := t - h
+	if n == 0 {
+		// Try runnext
+		next := pp.runnext
+		if next != unsafe { nil } {
+			if C.goroutines_atomic_cas_ptr(voidptr(&pp.runnext), voidptr(&next), unsafe { nil }) {
+				return next
+			}
+		}
+		return unsafe { nil }
+	}
+	// Steal half
+	steal := n - n / 2
+	mut first := unsafe { &Goroutine(nil) }
+	for i := u32(0); i < steal; i++ {
+		mut gp := pp.runq[(h + i) % local_queue_size]
+		if i == 0 {
+			first = gp
+		} else {
+			// Enqueue remaining stolen goroutines into thief's local run queue
+			runq_put(mut thisp, gp, false)
+		}
+	}
+	C.goroutines_atomic_fetch_add_u32(&pp.runq_head, steal)
+	return first
+}
+
+// Global run queue operations (translated from Go's globrunqput/get)
+fn glob_runq_put(gp &Goroutine) {
+	gsched.mu.acquire()
+	gsched.runq.push_back(gp)
+	gsched.mu.release()
+}
+
+fn glob_runq_get() &Goroutine {
+	gsched.mu.acquire()
+	gp := gsched.runq.pop()
+	gsched.mu.release()
+	return gp
+}
+
+// Local run queue operations (translated from Go's runqput/runqget)
+
+// runq_put puts a G on the local run queue.
+// If next is true, it goes into runnext for immediate scheduling.
+// Translated from Go's runqput() in proc.go.
+fn runq_put(mut pp Processor, gp &Goroutine, next bool) {
+	if next {
+		// Fast path: put as runnext
+		old := pp.runnext
+		pp.runnext = unsafe { gp }
+		if old == unsafe { nil } {
+			return
+		}
+		// Kick old runnext to the regular queue
+		runq_put(mut pp, old, false)
+		return
+	}
+
+	// Regular path: put on the ring buffer
+	h := C.goroutines_atomic_load_u32(&pp.runq_head)
+	t := pp.runq_tail
+	if t - h < local_queue_size {
+		pp.runq[t % local_queue_size] = unsafe { gp }
+		C.goroutines_atomic_store_u32(&pp.runq_tail, t + 1)
+		return
+	}
+	// Queue is full - put half on global queue
+	runq_put_slow(mut pp, gp, h, t)
+}
+
+// runq_put_slow moves half the local queue to the global queue.
+// Translated from Go's runqputslow() in proc.go.
+fn runq_put_slow(mut pp Processor, gp &Goroutine, h u32, t u32) {
+	n := (t - h) / 2
+	mut batch := GoroutineQueue{}
+	for i := u32(0); i < n; i++ {
+		g := pp.runq[(h + i) % local_queue_size]
+		batch.push_back(g)
+	}
+	C.goroutines_atomic_fetch_add_u32(&pp.runq_head, n)
+	batch.push_back(gp)
+
+	gsched.mu.acquire()
+	for !batch.empty() {
+		gsched.runq.push_back(batch.pop())
+	}
+	gsched.mu.release()
+}
+
+// runq_get gets a G from the local run queue.
+// Translated from Go's runqget() in proc.go.
+fn runq_get(mut pp Processor) (&Goroutine, bool) {
+	// Check runnext first (fast path)
+	next := pp.runnext
+	if next != unsafe { nil } {
+		pp.runnext = unsafe { nil }
+		return next, true
+	}
+
+	// Regular queue
+	for {
+		h := C.goroutines_atomic_load_u32(&pp.runq_head)
+		t := pp.runq_tail
+		if t == h {
+			return unsafe { nil }, false
+		}
+		gp := pp.runq[h % local_queue_size]
+		if C.goroutines_atomic_cas_u32(&pp.runq_head, &h, h + 1) {
+			return gp, false
+		}
+	}
+	return unsafe { nil }, false
+}
+
+// Idle P list operations
+fn pid_get() &Processor {
+	pp := gsched.pidle
+	if pp != unsafe { nil } {
+		gsched.pidle = pp.link
+		gsched.npidle--
+		mut p := unsafe { pp }
+		p.status = .running
+		p.link = unsafe { nil }
+	}
+	return pp
+}
+
+fn pid_put(pp &Processor) {
+	mut p := unsafe { pp }
+	p.status = .idle
+	p.m = unsafe { nil }
+	p.link = gsched.pidle
+	gsched.pidle = p
+	gsched.npidle++
+}
+
+// G free list operations (translated from Go's gfput/gfget)
+fn gfput(mut pp Processor, gp &Goroutine) {
+	mut g := unsafe { gp }
+	g.status = .dead
+	g.fn_ptr = unsafe { nil }
+	g.fn_arg = unsafe { nil }
+	g.sched_link = unsafe { nil }
+	pp.g_free.push(g)
+}
+
+fn gfget(mut pp Processor) &Goroutine {
+	return pp.g_free.pop()
+}
+
+fn gfput_global(gp &Goroutine) {
+	gsched.g_free_mu.acquire()
+	gsched.g_free.push(unsafe { gp })
+	gsched.g_free_count++
+	gsched.g_free_mu.release()
+}
+
+fn gfget_global() &Goroutine {
+	gsched.g_free_mu.acquire()
+	gp := gsched.g_free.pop()
+	if gp != unsafe { nil } {
+		gsched.g_free_count--
+	}
+	gsched.g_free_mu.release()
+	return gp
+}
+
+// assign_goid allocates a unique goroutine ID.
+// Uses per-P caching to avoid contention (like Go's goidcache).
+fn assign_goid() u64 {
+	mut pp := get_current_p()
+	if pp != unsafe { nil } && pp.goid_cache < pp.goid_cache_end {
+		id := pp.goid_cache
+		unsafe {
+			pp.goid_cache++
+		}
+		return id
+	}
+	// Refill cache from global counter
+	batch := u64(16)
+	id := C.goroutines_atomic_fetch_add_u64(&gsched.goid_gen, batch)
+	if pp != unsafe { nil } {
+		unsafe {
+			pp.goid_cache = id + 1
+			pp.goid_cache_end = id + batch
+		}
+	}
+	return id
+}
+
+// proc_yield spins for a short time (used during work stealing).
+fn proc_yield(count int) {
+	for _ in 0 .. count {
+		// CPU pause instruction to reduce power and contention
+		$if amd64 {
+			asm volatile amd64 {
+				pause
+			}
+		}
+		$if arm64 {
+			asm volatile arm64 {
+				yield
+			}
+		}
+	}
+}
+
+// get_current_m returns the M for the current OS thread.
+// Uses thread-local storage via C _Thread_local (see tls.c).
+fn get_current_m() &Machine {
+	return unsafe { &Machine(C.goroutines_get_current_m()) }
+}
+
+fn set_current_m(mp &Machine) {
+	C.goroutines_set_current_m(voidptr(mp))
+}
+
+// get_current_p returns the P for the current OS thread's M.
+fn get_current_p() &Processor {
+	mp := get_current_m()
+	if mp == unsafe { nil } {
+		return unsafe { nil }
+	}
+	return mp.p
+}
+
+// get_current_g returns the currently running G.
+pub fn get_current_g() &Goroutine {
+	mp := get_current_m()
+	if mp == unsafe { nil } {
+		return unsafe { nil }
+	}
+	return mp.curg
+}