From d0e75f2bef4e16356693987db6ae6bbdce749618 Mon Sep 17 00:00:00 2001
From: Ian Gudger <igudger@google.com>
Date: Tue, 21 Jan 2020 18:34:24 -0800
Subject: [PATCH] Add trylock support to DowngradableRWMutex.

Updates #231

PiperOrigin-RevId: 290868875
---
 pkg/sync/downgradable_rwmutex_test.go   | 55 ++++++++++++++++++++
 pkg/sync/downgradable_rwmutex_unsafe.go | 68 ++++++++++++++++++++++---
 2 files changed, 115 insertions(+), 8 deletions(-)

diff --git a/pkg/sync/downgradable_rwmutex_test.go b/pkg/sync/downgradable_rwmutex_test.go
index f04496bc5..b5cb28ec0 100644
--- a/pkg/sync/downgradable_rwmutex_test.go
+++ b/pkg/sync/downgradable_rwmutex_test.go
@@ -148,3 +148,58 @@ func TestDowngradableRWMutex(t *testing.T) {
 	HammerDowngradableRWMutex(10, 10, n)
 	HammerDowngradableRWMutex(10, 5, n)
 }
+
+func TestRWDoubleTryLock(t *testing.T) {
+	var m DowngradableRWMutex
+	if !m.TryLock() {
+		t.Fatal("failed to aquire lock")
+	}
+	if m.TryLock() {
+		t.Fatal("unexpectedly succeeded in aquiring locked mutex")
+	}
+}
+
+func TestRWTryLockAfterLock(t *testing.T) {
+	var m DowngradableRWMutex
+	m.Lock()
+	if m.TryLock() {
+		t.Fatal("unexpectedly succeeded in aquiring locked mutex")
+	}
+}
+
+func TestRWTryLockUnlock(t *testing.T) {
+	var m DowngradableRWMutex
+	if !m.TryLock() {
+		t.Fatal("failed to aquire lock")
+	}
+	m.Unlock()
+	if !m.TryLock() {
+		t.Fatal("failed to aquire lock after unlock")
+	}
+}
+
+func TestTryRLockAfterLock(t *testing.T) {
+	var m DowngradableRWMutex
+	m.Lock()
+	if m.TryRLock() {
+		t.Fatal("unexpectedly succeeded in aquiring locked mutex")
+	}
+}
+
+func TestTryLockAfterRLock(t *testing.T) {
+	var m DowngradableRWMutex
+	m.RLock()
+	if m.TryLock() {
+		t.Fatal("unexpectedly succeeded in aquiring locked mutex")
+	}
+}
+
+func TestDoubleTryRLock(t *testing.T) {
+	var m DowngradableRWMutex
+	if !m.TryRLock() {
+		t.Fatal("failed to aquire lock")
+	}
+	if !m.TryRLock() {
+		t.Fatal("failed to read aquire read locked lock")
+	}
+}
diff --git a/pkg/sync/downgradable_rwmutex_unsafe.go b/pkg/sync/downgradable_rwmutex_unsafe.go
index 9bb55cd3a..0d321f5e3 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,18 +29,43 @@ 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.
+// DowngradableRWMutex is identical to sync.RWMutex, but adds the DowngradeLock,
+// TryLock and TryRLock methods.
 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
+	w           TMutex // 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 *DowngradableRWMutex) 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() {
 	if RaceEnabled {
@@ -78,6 +102,34 @@ func (rw *DowngradableRWMutex) RUnlock() {
 	}
 }
 
+// TryLock locks rw for writing. It returns true if it succeeds and false
+// otherwise. It does not block.
+func (rw *DowngradableRWMutex) 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() {
 	if RaceEnabled {