syncs: add AssertLocked

This allows us to check lock invariants.

It was proposed upstream and rejected in:
https://github.com/golang/go/issues/1366

Signed-off-by: Josh Bleecher Snyder <josh@tailscale.com>

3 files changed, 182 insertions, 1 deletions

diff --git a/syncs/locked.go b/syncs/locked.go
new file mode 100644
index 000000000..8e8999076
--- /dev/null
+++ b/syncs/locked.go
@@ -0,0 +1,58 @@
+// Copyright (c) 2020 Tailscale Inc & AUTHORS All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// +build go1.13,!go1.16
+
+// This file makes assumptions about the inner workings of sync.Mutex and sync.RWMutex.
+// This includes not just their memory layout but their invariants and functionality.
+// To prevent accidents, it is limited to a known good subset of Go versions.
+
+package syncs
+
+import (
+	"sync"
+	"sync/atomic"
+	"unsafe"
+)
+
+const (
+	mutexLocked = 1
+
+	// sync.Mutex field offsets
+	stateOffset = 0
+
+	// sync.RWMutext field offsets
+	mutexOffset       = 0
+	readerCountOffset = 16
+)
+
+// add returns a pointer with value p + off.
+func add(p unsafe.Pointer, off uintptr) unsafe.Pointer {
+	return unsafe.Pointer(uintptr(p) + off)
+}
+
+// AssertLocked panics if m is not locked.
+func AssertLocked(m *sync.Mutex) {
+	p := add(unsafe.Pointer(m), stateOffset)
+	if atomic.LoadInt32((*int32)(p))&mutexLocked == 0 {
+		panic("mutex is not locked")
+	}
+}
+
+// AssertRLocked panics if rw is not locked for reading or writing.
+func AssertRLocked(rw *sync.RWMutex) {
+	p := add(unsafe.Pointer(rw), readerCountOffset)
+	if atomic.LoadInt32((*int32)(p)) != 0 {
+		// There are readers present or writers pending, so someone has a read lock.
+		return
+	}
+	// No readers.
+	AssertWLocked(rw)
+}
+
+// AssertWLocked panics if rw is not locked for writing.
+func AssertWLocked(rw *sync.RWMutex) {
+	m := (*sync.Mutex)(add(unsafe.Pointer(rw), mutexOffset))
+	AssertLocked(m)
+}
diff --git a/syncs/locked_test.go b/syncs/locked_test.go
new file mode 100644
index 000000000..32b9b64e0
--- /dev/null
+++ b/syncs/locked_test.go
@@ -0,0 +1,123 @@
+// Copyright (c) 2020 Tailscale Inc & AUTHORS All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// +build go1.13,!go1.16
+
+//lint:file-ignore SA2001 the empty critical sections are part of triggering different internal mutex states
+
+package syncs
+
+import (
+	"sync"
+	"testing"
+	"time"
+)
+
+func wantPanic(t *testing.T, fn func()) {
+	t.Helper()
+	defer func() {
+		recover()
+	}()
+	fn()
+	t.Fatal("failed to panic")
+}
+
+func TestAssertLocked(t *testing.T) {
+	m := new(sync.Mutex)
+	wantPanic(t, func() { AssertLocked(m) })
+	m.Lock()
+	AssertLocked(m)
+	m.Unlock()
+	wantPanic(t, func() { AssertLocked(m) })
+	// Test correct handling of mutex with waiter.
+	m.Lock()
+	AssertLocked(m)
+	go func() {
+		m.Lock()
+		m.Unlock()
+	}()
+	// Give the goroutine above a few moments to get started.
+	// The test will pass whether or not we win the race,
+	// but we want to run sometimes, to get the test coverage.
+	time.Sleep(10 * time.Millisecond)
+	AssertLocked(m)
+}
+
+func TestAssertWLocked(t *testing.T) {
+	m := new(sync.RWMutex)
+	wantPanic(t, func() { AssertWLocked(m) })
+	m.Lock()
+	AssertWLocked(m)
+	m.Unlock()
+	wantPanic(t, func() { AssertWLocked(m) })
+	// Test correct handling of mutex with waiter.
+	m.Lock()
+	AssertWLocked(m)
+	go func() {
+		m.Lock()
+		m.Unlock()
+	}()
+	// Give the goroutine above a few moments to get started.
+	// The test will pass whether or not we win the race,
+	// but we want to run sometimes, to get the test coverage.
+	time.Sleep(10 * time.Millisecond)
+	AssertWLocked(m)
+}
+
+func TestAssertRLocked(t *testing.T) {
+	m := new(sync.RWMutex)
+	wantPanic(t, func() { AssertRLocked(m) })
+
+	m.Lock()
+	AssertRLocked(m)
+	m.Unlock()
+
+	m.RLock()
+	AssertRLocked(m)
+	m.RUnlock()
+
+	wantPanic(t, func() { AssertRLocked(m) })
+
+	// Test correct handling of mutex with waiter.
+	m.RLock()
+	AssertRLocked(m)
+	go func() {
+		m.RLock()
+		m.RUnlock()
+	}()
+	// Give the goroutine above a few moments to get started.
+	// The test will pass whether or not we win the race,
+	// but we want to run sometimes, to get the test coverage.
+	time.Sleep(10 * time.Millisecond)
+	AssertRLocked(m)
+	m.RUnlock()
+
+	// Test correct handling of rlock with write waiter.
+	m.RLock()
+	AssertRLocked(m)
+	go func() {
+		m.Lock()
+		m.Unlock()
+	}()
+	// Give the goroutine above a few moments to get started.
+	// The test will pass whether or not we win the race,
+	// but we want to run sometimes, to get the test coverage.
+	time.Sleep(10 * time.Millisecond)
+	AssertRLocked(m)
+	m.RUnlock()
+
+	// Test correct handling of rlock with other rlocks.
+	// This is a bit racy, but losing the race hurts nothing,
+	// and winning the race means correct test coverage.
+	m.RLock()
+	AssertRLocked(m)
+	go func() {
+		m.RLock()
+		time.Sleep(10 * time.Millisecond)
+		m.RUnlock()
+	}()
+	time.Sleep(5 * time.Millisecond)
+	AssertRLocked(m)
+	m.RUnlock()
+}
diff --git a/syncs/syncs.go b/syncs/syncs.go
index 4ea232335..c0208c996 100644
--- a/syncs/syncs.go
+++ b/syncs/syncs.go
@@ -2,7 +2,7 @@
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
 
-// Package syncs contains addition sync types.
+// Package syncs contains additional sync types and functionality.
 package syncs
 
 import "sync/atomic"