From cb396113893bbd1f03c19573a849bf424e7a7ca1 Mon Sep 17 00:00:00 2001
From: soypat <graded.sp@gmail.com>
Date: Wed, 1 Nov 2023 21:16:26 -0300
Subject: [PATCH] sync: implement trylock

---
 src/sync/mutex.go      | 40 +++++++++++++++++++++++++++++++++++-----
 src/sync/mutex_test.go | 36 ++++++++++++++++++++++++++++++++++++
 2 files changed, 71 insertions(+), 5 deletions(-)

diff --git a/src/sync/mutex.go b/src/sync/mutex.go
index e12bf40c..59f320d5 100644
--- a/src/sync/mutex.go
+++ b/src/sync/mutex.go
@@ -3,10 +3,12 @@ package sync
 import (
 	"internal/task"
 	_ "unsafe"
+
+	"runtime/volatile"
 )
 
 type Mutex struct {
-	locked  bool
+	state   uint8 // Set to non-zero if locked.
 	blocked task.Stack
 }
 
@@ -14,18 +16,18 @@ type Mutex struct {
 func scheduleTask(*task.Task)
 
 func (m *Mutex) Lock() {
-	if m.locked {
+	if m.islocked() {
 		// Push self onto stack of blocked tasks, and wait to be resumed.
 		m.blocked.Push(task.Current())
 		task.Pause()
 		return
 	}
 
-	m.locked = true
+	m.setlock(true)
 }
 
 func (m *Mutex) Unlock() {
-	if !m.locked {
+	if !m.islocked() {
 		panic("sync: unlock of unlocked Mutex")
 	}
 
@@ -33,8 +35,36 @@ func (m *Mutex) Unlock() {
 	if t := m.blocked.Pop(); t != nil {
 		scheduleTask(t)
 	} else {
-		m.locked = false
+		m.setlock(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.islocked() {
+		return false
+	}
+	m.Lock()
+	return true
+}
+
+func (m *Mutex) islocked() bool {
+	return volatile.LoadUint8(&m.state) != 0
+}
+
+func (m *Mutex) setlock(b bool) {
+	volatile.StoreUint8(&m.state, boolToU8(b))
+}
+
+func boolToU8(b bool) uint8 {
+	if b {
+		return 1
 	}
+	return 0
 }
 
 type RWMutex struct {
diff --git a/src/sync/mutex_test.go b/src/sync/mutex_test.go
index 88ae317d..accb01c9 100644
--- a/src/sync/mutex_test.go
+++ b/src/sync/mutex_test.go
@@ -7,6 +7,42 @@ import (
 	"testing"
 )
 
+func HammerMutex(m *sync.Mutex, loops int, cdone chan bool) {
+	for i := 0; i < loops; i++ {
+		if i%3 == 0 {
+			if m.TryLock() {
+				m.Unlock()
+			}
+			continue
+		}
+		m.Lock()
+		m.Unlock()
+	}
+	cdone <- true
+}
+
+func TestMutex(t *testing.T) {
+	m := new(sync.Mutex)
+
+	m.Lock()
+	if m.TryLock() {
+		t.Fatalf("TryLock succeeded with mutex locked")
+	}
+	m.Unlock()
+	if !m.TryLock() {
+		t.Fatalf("TryLock failed with mutex unlocked")
+	}
+	m.Unlock()
+
+	c := make(chan bool)
+	for i := 0; i < 10; i++ {
+		go HammerMutex(m, 1000, c)
+	}
+	for i := 0; i < 10; i++ {
+		<-c
+	}
+}
+
 // TestMutexUncontended tests locking and unlocking a Mutex that is not shared with any other goroutines.
 func TestMutexUncontended(t *testing.T) {
 	var mu sync.Mutex