From 89dd039cae865b32b1e29bb1a1489b5493eb9483 Mon Sep 17 00:00:00 2001
From: Ayke van Laethem <aykevanlaethem@gmail.com>
Date: Mon, 28 Oct 2024 15:10:04 +0100
Subject: [PATCH] sync: implement futex-based Mutex

This is largely based on this paper by Ulrich Drepper:
https://dept-info.labri.fr/~denis/Enseignement/2008-IR/Articles/01-futex.pdf

See the comment in the source code for details.
---
 src/internal/task/mutex.go | 69 ++++++++++++++++++++++++++++++++++++++
 src/sync/cond.go           |  3 ++
 src/sync/mutex.go          | 45 +------------------------
 3 files changed, 73 insertions(+), 44 deletions(-)
 create mode 100644 src/internal/task/mutex.go

diff --git a/src/internal/task/mutex.go b/src/internal/task/mutex.go
new file mode 100644
index 0000000000..97087de69f
--- /dev/null
+++ b/src/internal/task/mutex.go
@@ -0,0 +1,69 @@
+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
+	// look. 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/sync/cond.go b/src/sync/cond.go
index fb5f224927..139d8e0229 100644
--- a/src/sync/cond.go
+++ b/src/sync/cond.go
@@ -89,3 +89,6 @@ func (c *Cond) Wait() {
 	// signal.
 	task.Pause()
 }
+
+//go:linkname scheduleTask runtime.scheduleTask
+func scheduleTask(*task.Task)
diff --git a/src/sync/mutex.go b/src/sync/mutex.go
index b62b9fafdb..770735b3d9 100644
--- a/src/sync/mutex.go
+++ b/src/sync/mutex.go
@@ -5,50 +5,7 @@ import (
 	_ "unsafe"
 )
 
-type Mutex struct {
-	locked  bool
-	blocked task.Stack
-}
-
-//go:linkname scheduleTask runtime.scheduleTask
-func scheduleTask(*task.Task)
-
-func (m *Mutex) Lock() {
-	if m.locked {
-		// Push self onto stack of blocked tasks, and wait to be resumed.
-		m.blocked.Push(task.Current())
-		task.Pause()
-		return
-	}
-
-	m.locked = true
-}
-
-func (m *Mutex) Unlock() {
-	if !m.locked {
-		panic("sync: unlock of unlocked Mutex")
-	}
-
-	// Wake up a blocked task, if applicable.
-	if t := m.blocked.Pop(); t != nil {
-		scheduleTask(t)
-	} else {
-		m.locked = false
-	}
-}
-
-// 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 {
-	if m.locked {
-		return false
-	}
-	m.Lock()
-	return true
-}
+type Mutex = task.Mutex
 
 type RWMutex struct {
 	// waitingWriters are all of the tasks waiting for write locks.