1 files changed, 69 insertions, 17 deletions

diff --git a/pkg/sync/downgradable_rwmutex_unsafe.go b/pkg/sync/downgradable_rwmutex_unsafe.go
index 9bb55cd3a..ea6cdc447 100644
--- a/pkg/sync/downgradable_rwmutex_unsafe.go
+++ b/pkg/sync/downgradable_rwmutex_unsafe.go
@@ -19,7 +19,6 @@
 package sync
 
 import (
-	"sync"
 	"sync/atomic"
 	"unsafe"
 )
@@ -30,20 +29,45 @@ func runtimeSemacquire(s *uint32)
 //go:linkname runtimeSemrelease sync.runtime_Semrelease
 func runtimeSemrelease(s *uint32, handoff bool, skipframes int)
 
-// DowngradableRWMutex is identical to sync.RWMutex, but adds the DowngradeLock
-// method.
-type DowngradableRWMutex struct {
-	w           sync.Mutex // held if there are pending writers
-	writerSem   uint32     // semaphore for writers to wait for completing readers
-	readerSem   uint32     // semaphore for readers to wait for completing writers
-	readerCount int32      // number of pending readers
-	readerWait  int32      // number of departing readers
+// RWMutex is identical to sync.RWMutex, but adds the DowngradeLock,
+// TryLock and TryRLock methods.
+type RWMutex struct {
+	w           Mutex  // held if there are pending writers
+	writerSem   uint32 // semaphore for writers to wait for completing readers
+	readerSem   uint32 // semaphore for readers to wait for completing writers
+	readerCount int32  // number of pending readers
+	readerWait  int32  // number of departing readers
 }
 
 const rwmutexMaxReaders = 1 << 30
 
+// TryRLock locks rw for reading. It returns true if it succeeds and false
+// otherwise. It does not block.
+func (rw *RWMutex) TryRLock() bool {
+	if RaceEnabled {
+		RaceDisable()
+	}
+	for {
+		rc := atomic.LoadInt32(&rw.readerCount)
+		if rc < 0 {
+			if RaceEnabled {
+				RaceEnable()
+			}
+			return false
+		}
+		if !atomic.CompareAndSwapInt32(&rw.readerCount, rc, rc+1) {
+			continue
+		}
+		if RaceEnabled {
+			RaceEnable()
+			RaceAcquire(unsafe.Pointer(&rw.readerSem))
+		}
+		return true
+	}
+}
+
 // RLock locks rw for reading.
-func (rw *DowngradableRWMutex) RLock() {
+func (rw *RWMutex) RLock() {
 	if RaceEnabled {
 		RaceDisable()
 	}
@@ -58,14 +82,14 @@ func (rw *DowngradableRWMutex) RLock() {
 }
 
 // RUnlock undoes a single RLock call.
-func (rw *DowngradableRWMutex) RUnlock() {
+func (rw *RWMutex) RUnlock() {
 	if RaceEnabled {
 		RaceReleaseMerge(unsafe.Pointer(&rw.writerSem))
 		RaceDisable()
 	}
 	if r := atomic.AddInt32(&rw.readerCount, -1); r < 0 {
 		if r+1 == 0 || r+1 == -rwmutexMaxReaders {
-			panic("RUnlock of unlocked DowngradableRWMutex")
+			panic("RUnlock of unlocked RWMutex")
 		}
 		// A writer is pending.
 		if atomic.AddInt32(&rw.readerWait, -1) == 0 {
@@ -78,8 +102,36 @@ func (rw *DowngradableRWMutex) RUnlock() {
 	}
 }
 
+// TryLock locks rw for writing. It returns true if it succeeds and false
+// otherwise. It does not block.
+func (rw *RWMutex) TryLock() bool {
+	if RaceEnabled {
+		RaceDisable()
+	}
+	// First, resolve competition with other writers.
+	if !rw.w.TryLock() {
+		if RaceEnabled {
+			RaceEnable()
+		}
+		return false
+	}
+	// Only proceed if there are no readers.
+	if !atomic.CompareAndSwapInt32(&rw.readerCount, 0, -rwmutexMaxReaders) {
+		rw.w.Unlock()
+		if RaceEnabled {
+			RaceEnable()
+		}
+		return false
+	}
+	if RaceEnabled {
+		RaceEnable()
+		RaceAcquire(unsafe.Pointer(&rw.writerSem))
+	}
+	return true
+}
+
 // Lock locks rw for writing.
-func (rw *DowngradableRWMutex) Lock() {
+func (rw *RWMutex) Lock() {
 	if RaceEnabled {
 		RaceDisable()
 	}
@@ -98,7 +150,7 @@ func (rw *DowngradableRWMutex) Lock() {
 }
 
 // Unlock unlocks rw for writing.
-func (rw *DowngradableRWMutex) Unlock() {
+func (rw *RWMutex) Unlock() {
 	if RaceEnabled {
 		RaceRelease(unsafe.Pointer(&rw.writerSem))
 		RaceRelease(unsafe.Pointer(&rw.readerSem))
@@ -107,7 +159,7 @@ func (rw *DowngradableRWMutex) Unlock() {
 	// Announce to readers there is no active writer.
 	r := atomic.AddInt32(&rw.readerCount, rwmutexMaxReaders)
 	if r >= rwmutexMaxReaders {
-		panic("Unlock of unlocked DowngradableRWMutex")
+		panic("Unlock of unlocked RWMutex")
 	}
 	// Unblock blocked readers, if any.
 	for i := 0; i < int(r); i++ {
@@ -121,7 +173,7 @@ func (rw *DowngradableRWMutex) Unlock() {
 }
 
 // DowngradeLock atomically unlocks rw for writing and locks it for reading.
-func (rw *DowngradableRWMutex) DowngradeLock() {
+func (rw *RWMutex) DowngradeLock() {
 	if RaceEnabled {
 		RaceRelease(unsafe.Pointer(&rw.readerSem))
 		RaceDisable()
@@ -129,7 +181,7 @@ func (rw *DowngradableRWMutex) DowngradeLock() {
 	// Announce to readers there is no active writer and one additional reader.
 	r := atomic.AddInt32(&rw.readerCount, rwmutexMaxReaders+1)
 	if r >= rwmutexMaxReaders+1 {
-		panic("DowngradeLock of unlocked DowngradableRWMutex")
+		panic("DowngradeLock of unlocked RWMutex")
 	}
 	// Unblock blocked readers, if any. Note that this loop starts as 1 since r
 	// includes this goroutine.