From 5b243f652cb9bbc3629db3599735cd684c10830a Mon Sep 17 00:00:00 2001
From: Ayke van Laethem <aykevanlaethem@gmail.com>
Date: Sun, 27 Oct 2024 10:26:23 +0100
Subject: [PATCH] internal/task: implement atomic primitives for preemptive
 scheduling

---
 src/internal/task/atomic-cooperative.go |  2 +
 src/internal/task/atomic-preemptive.go  | 14 +++++
 src/internal/task/futex-cooperative.go  |  2 +
 src/internal/task/futex-preemptive.go   |  7 +++
 src/internal/task/mutex-cooperative.go  |  2 +
 src/internal/task/mutex-preemptive.go   | 71 +++++++++++++++++++++++++
 src/internal/task/pmutex-cooperative.go |  2 +
 src/internal/task/pmutex-preemptive.go  | 11 ++++
 8 files changed, 111 insertions(+)
 create mode 100644 src/internal/task/atomic-preemptive.go
 create mode 100644 src/internal/task/futex-preemptive.go
 create mode 100644 src/internal/task/mutex-preemptive.go
 create mode 100644 src/internal/task/pmutex-preemptive.go

diff --git a/src/internal/task/atomic-cooperative.go b/src/internal/task/atomic-cooperative.go
index 60eb917a8..bd4cba895 100644
--- a/src/internal/task/atomic-cooperative.go
+++ b/src/internal/task/atomic-cooperative.go
@@ -1,3 +1,5 @@
+//go:build !scheduler.threads
+
 package task
 
 // Atomics implementation for cooperative systems. The atomic types here aren't
diff --git a/src/internal/task/atomic-preemptive.go b/src/internal/task/atomic-preemptive.go
new file mode 100644
index 000000000..275f36dce
--- /dev/null
+++ b/src/internal/task/atomic-preemptive.go
@@ -0,0 +1,14 @@
+//go:build scheduler.threads
+
+package task
+
+// Atomics implementation for non-cooperative systems (multithreaded, etc).
+// These atomic types use real atomic instructions.
+
+import "sync/atomic"
+
+type (
+	Uintptr = atomic.Uintptr
+	Uint32  = atomic.Uint32
+	Uint64  = atomic.Uint64
+)
diff --git a/src/internal/task/futex-cooperative.go b/src/internal/task/futex-cooperative.go
index 8351f8877..2a42c28d4 100644
--- a/src/internal/task/futex-cooperative.go
+++ b/src/internal/task/futex-cooperative.go
@@ -1,3 +1,5 @@
+//go:build !scheduler.threads
+
 package task
 
 // A futex is a way for userspace to wait with the pointer as the key, and for
diff --git a/src/internal/task/futex-preemptive.go b/src/internal/task/futex-preemptive.go
new file mode 100644
index 000000000..7f9e89580
--- /dev/null
+++ b/src/internal/task/futex-preemptive.go
@@ -0,0 +1,7 @@
+//go:build scheduler.threads
+
+package task
+
+import "internal/futex"
+
+type Futex = futex.Futex
diff --git a/src/internal/task/mutex-cooperative.go b/src/internal/task/mutex-cooperative.go
index e40966bed..f1205eea2 100644
--- a/src/internal/task/mutex-cooperative.go
+++ b/src/internal/task/mutex-cooperative.go
@@ -1,3 +1,5 @@
+//go:build !scheduler.threads
+
 package task
 
 type Mutex struct {
diff --git a/src/internal/task/mutex-preemptive.go b/src/internal/task/mutex-preemptive.go
new file mode 100644
index 000000000..27f464669
--- /dev/null
+++ b/src/internal/task/mutex-preemptive.go
@@ -0,0 +1,71 @@
+//go:build scheduler.threads
+
+package task
+
+// Futex-based mutex.
+// This is largely based on the paper "Futexes are Tricky" by Ulrich Drepper.
+// It describes a few ways to implement mutexes using a futex, and how some
+// seemingly-obvious implementations don't exactly work as intended.
+// Unfortunately, Go atomic operations work slightly differently so we can't
+// copy the algorithm verbatim.
+//
+// The implementation works like this. The futex can have 3 different values,
+// depending on the state:
+//
+//   - 0: the futex is currently unlocked.
+//   - 1: the futex is locked, but is uncontended. There is one special case: if
+//     a contended futex is unlocked, it is set to 0. It is possible for another
+//     thread to lock the futex before the next waiter is woken. But because a
+//     waiter will be woken (if there is one), it will always change to 2
+//     regardless. So this is not a problem.
+//   - 2: the futex is locked, and is contended. At least one thread is trying
+//     to obtain the lock (and is in the contended loop, see below).
+//
+// For the paper, see:
+// https://dept-info.labri.fr/~denis/Enseignement/2008-IR/Articles/01-futex.pdf)
+
+type Mutex struct {
+	futex Futex
+}
+
+func (m *Mutex) Lock() {
+	// Fast path: try to take an uncontended lock.
+	if m.futex.CompareAndSwap(0, 1) {
+		// We obtained the mutex.
+		return
+	}
+
+	// The futex is contended, so we enter the contended loop.
+	// If we manage to change the futex from 0 to 2, we managed to take the
+	// lock. Else, we have to wait until a call to Unlock unlocks this mutex.
+	// (Unlock will wake one waiter when it finds the futex is set to 2 when
+	// unlocking).
+	for m.futex.Swap(2) != 0 {
+		// Wait until we get resumed in Unlock.
+		m.futex.Wait(2)
+	}
+}
+
+func (m *Mutex) Unlock() {
+	if old := m.futex.Swap(0); old == 0 {
+		// Mutex wasn't locked before.
+		panic("sync: unlock of unlocked Mutex")
+	} else if old == 2 {
+		// Mutex was a contended lock, so we need to wake the next waiter.
+		m.futex.Wake()
+	}
+}
+
+// TryLock tries to lock m and reports whether it succeeded.
+//
+// Note that while correct uses of TryLock do exist, they are rare,
+// and use of TryLock is often a sign of a deeper problem
+// in a particular use of mutexes.
+func (m *Mutex) TryLock() bool {
+	// Fast path: try to take an uncontended lock.
+	if m.futex.CompareAndSwap(0, 1) {
+		// We obtained the mutex.
+		return true
+	}
+	return false
+}
diff --git a/src/internal/task/pmutex-cooperative.go b/src/internal/task/pmutex-cooperative.go
index ae2aa4bad..0e6c4f828 100644
--- a/src/internal/task/pmutex-cooperative.go
+++ b/src/internal/task/pmutex-cooperative.go
@@ -1,3 +1,5 @@
+//go:build !scheduler.threads
+
 package task
 
 // PMutex is a real mutex on systems that can be either preemptive or threaded,
diff --git a/src/internal/task/pmutex-preemptive.go b/src/internal/task/pmutex-preemptive.go
new file mode 100644
index 000000000..10f0a6356
--- /dev/null
+++ b/src/internal/task/pmutex-preemptive.go
@@ -0,0 +1,11 @@
+//go:build scheduler.threads
+
+package task
+
+// PMutex is a real mutex on systems that can be either preemptive or threaded,
+// and a dummy lock on other (purely cooperative) systems.
+//
+// It is mainly useful for short operations that need a lock when threading may
+// be involved, but which do not need a lock with a purely cooperative
+// scheduler.
+type PMutex = Mutex