1 files changed, 212 insertions, 60 deletions
diff --git a/util/ctxlock/doc_test.go b/util/ctxlock/doc_test.go
index a6b0de407..d1c32f9ba 100644
--- a/util/ctxlock/doc_test.go
+++ b/util/ctxlock/doc_test.go
@@ -6,75 +6,130 @@ package ctxlock_test
 import (
 	"context"
 	"fmt"
-	"sync"
+	"strings"
 	"testing"
 
 	"tailscale.com/util/ctxlock"
 )
 
-type Resource struct {
-	mu       sync.Mutex
-	foo, bar string
-}
+func ExampleMutex_reentrant() {
+	var mu ctxlock.ReentrantMutex // shorthand for ctxlock.Mutex[ctxlock.Reentrant]
 
-func (r *Resource) GetFoo(ctx ctxlock.State) string {
-	// Lock the mutex if not already held.
-	defer ctxlock.Lock(ctx, &r.mu).Unlock()
-	return r.foo
-}
+	// The mutex is reentrant, so foo can be called with or without holding the mu.
+	// If mu is not already held, it will be locked on entry and unlocked on exit.
+	// The [ctxlock.State] parameter carries the current lock state.
+	foo := func(ctx ctxlock.State, msg string) {
+		lock := ctxlock.Lock(ctx, &mu)
+		defer lock.Unlock()
+		fmt.Println(msg)
+	}
 
-func (r *Resource) SetFoo(ctx ctxlock.State, foo string) {
-	// You can do it this way, if you prefer
-	// or if you need to pass the state to another function.
-	ctx = ctxlock.Lock(ctx, &r.mu)
-	defer ctx.Unlock()
-	r.foo = foo
-}
+	// Calling foo without holding the lock.
+	foo(ctxlock.None(), "no lock")
 
-func (r *Resource) GetBar(ctx ctxlock.State) string {
-	defer ctxlock.Lock(ctx, &r.mu).Unlock()
-	return r.bar
-}
+	// Locking the mutex and calling foo again.
+	lock := ctxlock.Lock(ctxlock.None(), &mu)
+	foo(lock.State(), "with lock")
+	defer lock.Unlock()
 
-func (r *Resource) SetBar(ctx ctxlock.State, bar string) {
-	defer ctxlock.Lock(ctx, &r.mu).Unlock()
-	r.bar = bar
+	// Output:
+	// no lock
+	// with lock
 }
 
-func (r *Resource) WithLock(ctx ctxlock.State, f func(ctx ctxlock.State)) {
-	// Lock the mutex if not already held, and get a new state.
-	ctx = ctxlock.Lock(ctx, &r.mu)
-	defer ctx.Unlock()
-	f(ctx) // Call the callback with the new lock state.
+func ExampleMutex_nonReentrant() {
+	var mu ctxlock.Mutex[ctxlock.NonReentrant]
+
+	// The mutex is non-reentrant, so foo must only be called without holding the mu.
+	// If mu is already held, it will panic attempting to lock it again.
+	foo := func(ctx ctxlock.State, msg string) {
+		defer func() {
+			if r := recover(); r != nil {
+				fmt.Println("panic:", trimPanicMessage(r))
+			}
+		}()
+
+		lock := ctxlock.Lock(ctx, &mu)
+		defer lock.Unlock()
+		fmt.Println(msg)
+	}
+
+	// Calling foo without holding the lock.
+	foo(ctxlock.None(), "no lock")
+
+	// Locking the mutex and calling foo again.
+	// This will panic because the mutex is non-reentrant.
+	lock := ctxlock.Lock(ctxlock.None(), &mu)
+	foo(lock.State(), "with lock")
+	defer lock.Unlock()
+
+	// Output:
+	// no lock
+	// panic: non-reentrant mutex already locked
 }
 
-func (r *Resource) HandleRequest(ctx context.Context, foo, bar string, f func(ls ctxlock.State) string) string {
-	// Same, but with a standard [context.Context] instead of [ctxlock.State].
-	// [ctxlock.Lock] is generic and works with both without allocating.
-	// The ctx can be used for cancellation, etc.
-	mu := ctxlock.Lock(ctx, &r.mu)
-	defer mu.Unlock()
-	r.foo = foo
-	r.bar = bar
-	return f(mu)
+func ExampleRank() {
+	var mu1 ctxlock.Mutex[rank1] // cannot be locked after mu2 or mu3
+	var mu2 ctxlock.Mutex[rank2] // cannot be locked after mu3
+	var mu3 ctxlock.Mutex[rank3]
+
+	lock := ctxlock.Lock(ctxlock.None(), &mu1)
+	defer lock.Unlock()
+	fmt.Println("locked mu1")
+
+	lock = ctxlock.Lock(lock.State(), &mu2)
+	defer lock.Unlock()
+	fmt.Println("locked mu2")
+
+	lock = ctxlock.Lock(lock.State(), &mu3)
+	defer lock.Unlock()
+	fmt.Println("locked mu3")
+
+	// Output:
+	// locked mu1
+	// locked mu2
+	// locked mu3
 }
 
-func (r *Resource) HandleIntRequest(ctx context.Context, foo, bar string, f func(ls ctxlock.State) int) int {
-	// Same, but returns an int instead of a string,
-	// and must not allocate with the unchecked implementation.
-	mu := ctxlock.Lock(ctx, &r.mu)
-	defer mu.Unlock()
-	r.foo = foo
-	r.bar = bar
-	return f(mu)
+func ExampleRank_lockOrderViolation() {
+	var mu1 ctxlock.Mutex[rank1] // cannot be locked after mu2 or mu3
+	var mu2 ctxlock.Mutex[rank2] // cannot be locked after mu3
+	var mu3 ctxlock.Mutex[rank3]
+
+	defer func() {
+		if r := recover(); r != nil {
+			fmt.Println("panic:", trimPanicMessage(r))
+		}
+	}()
+
+	// While we can lock mu2 first...
+	lock := ctxlock.Lock(ctxlock.None(), &mu2)
+	defer lock.Unlock()
+	fmt.Println("locked mu2")
+
+	// ...and then mu3...
+	lock = ctxlock.Lock(lock.State(), &mu3)
+	defer lock.Unlock()
+	fmt.Println("locked mu3")
+
+	// It is a lock order violation to lock mu1
+	// after either mu2 or mu3.
+	lock = ctxlock.Lock(lock.State(), &mu1)
+	defer lock.Unlock()
+	fmt.Println("locked mu1")
+
+	// Output:
+	// locked mu2
+	// locked mu3
+	// panic: cannot lock ctxlock_test.rank1 after ctxlock_test.rank3
 }
 
-func ExampleState() {
+func ExampleState_resource() {
 	var r Resource
 	r.SetFoo(ctxlock.None(), "foo")
 	r.SetBar(ctxlock.None(), "bar")
 	r.WithLock(ctxlock.None(), func(ctx ctxlock.State) {
-		// This callback is invoked with r's lock held,
+		// This callback is invoked with r's mutex held,
 		// and ctx carries the lock state. This means we can safely call
 		// other methods on r using ctx without causing a deadlock.
 		r.SetFoo(ctx, r.GetFoo(ctx)+r.GetBar(ctx))
@@ -88,7 +143,7 @@ func ExampleState_twoResources() {
 	r1.SetFoo(ctxlock.None(), "foo")
 	r2.SetBar(ctxlock.None(), "bar")
 	r1.WithLock(ctxlock.None(), func(ctx ctxlock.State) {
-		// Here, r1's lock is held, but r2's lock is not.
+		// Here, r1's mutex is held, but r2's mutex is not.
 		// So r2 will be locked when we call r2.GetBar(ctx).
 		r1.SetFoo(ctx, r1.GetFoo(ctx)+r2.GetBar(ctx))
 	})
@@ -96,29 +151,27 @@ func ExampleState_twoResources() {
 	// Output: foobar
 }
 
-func ExampleState_stdContext() {
+func ExampleState_withStdContext() {
 	var r Resource
 	ctx := context.Background()
 	result := r.HandleRequest(ctx, "foo", "bar", func(ctx ctxlock.State) string {
-		// The r's lock is held, and ctx carries the lock state.
+		// The r's mutex is held, and ctx carries the lock state.
 		return r.GetFoo(ctx) + r.GetBar(ctx)
 	})
 	fmt.Println(result)
 	// Output: foobar
 }
 
-func TestAllocFree(t *testing.T) {
-	if ctxlock.Checked {
-		t.Skip("Exported implementation is not alloc-free (use --tags=ts_omit_ctxlock_checks)")
+func TestEndToEndAllocFree(t *testing.T) {
+	if ctxlock.IsChecked {
+		t.Skip("Exported implementation is not alloc-free (use --tags=ts_omit_ctxlock_checks).")
 	}
 
 	var r Resource
-	ctx := context.Background()
-
-	const runs = 1000
-	if allocs := testing.AllocsPerRun(runs, func() {
-		res := r.HandleIntRequest(ctx, "foo", "bar", func(ctx ctxlock.State) int {
-			// The r's lock is held, and ctx carries the lock state.
+	const N = 1000
+	if allocs := testing.AllocsPerRun(N, func() {
+		res := r.HandleIntRequest(context.Background(), "foo", "bar", func(ctx ctxlock.State) int {
+			// The r's mutex is held, and ctx carries the lock state.
 			return len(r.GetFoo(ctx) + r.GetBar(ctx))
 		})
 		if res != 6 {
@@ -128,3 +181,102 @@ func TestAllocFree(t *testing.T) {
 		t.Errorf("expected 0 allocs, got %f", allocs)
 	}
 }
+
+type (
+	rank1 struct{}
+	rank2 struct{}
+	rank3 struct{}
+)
+
+// CheckLockAfter implements [ctxlock.Rank].
+func (r rank1) CheckLockAfter(r2 ctxlock.Rank) error {
+	switch r2.(type) {
+	case rank2, rank3:
+		return fmt.Errorf("cannot lock %T after %T", r, r2)
+	default:
+		return nil
+	}
+}
+
+// CheckLockAfter implements [ctxlock.Rank].
+func (r rank2) CheckLockAfter(r2 ctxlock.Rank) error {
+	switch r2.(type) {
+	case rank2, rank3:
+		return fmt.Errorf("cannot lock %T after %T", r, r2)
+	default:
+		return nil
+	}
+}
+
+// CheckLockAfter implements [ctxlock.Rank].
+func (a rank3) CheckLockAfter(b ctxlock.Rank) error {
+	return nil
+}
+
+type Resource struct {
+	mu       ctxlock.ReentrantMutex
+	foo, bar string
+}
+
+func (r *Resource) GetFoo(ctx ctxlock.State) string {
+	// Lock the mutex if not already held,
+	// and unlock it when the function returns.
+	defer ctxlock.Lock(ctx, &r.mu).Unlock()
+	return r.foo
+}
+
+func (r *Resource) SetFoo(ctx ctxlock.State, foo string) {
+	// You can do it this way, if you prefer.
+	mu := ctxlock.Lock(ctx, &r.mu)
+	defer mu.Unlock()
+	r.foo = foo
+}
+
+func (r *Resource) GetBar(ctx ctxlock.State) string {
+	mu := ctxlock.Lock(ctx, &r.mu)
+	defer mu.Unlock()
+	return r.bar
+}
+
+func (r *Resource) SetBar(ctx ctxlock.State, bar string) {
+	mu := ctxlock.Lock(ctx, &r.mu)
+	defer mu.Unlock()
+	r.bar = bar
+}
+
+func (r *Resource) WithLock(ctx ctxlock.State, f func(ctx ctxlock.State)) {
+	mu := ctxlock.Lock(ctx, &r.mu)
+	defer mu.Unlock()
+	// Call the callback with the new lock state.
+	f(mu.State())
+}
+
+func (r *Resource) HandleRequest(ctx context.Context, foo, bar string, f func(ls ctxlock.State) string) string {
+	// Same, but with a standard [context.Context] instead of [ctxlock.State].
+	// [ctxlock.Lock] is generic and works with both without allocating.
+	// The ctx can be used for cancellation, etc.
+	mu := ctxlock.Lock(ctx, &r.mu)
+	defer mu.Unlock()
+	r.foo = foo
+	r.bar = bar
+	return f(mu.State())
+}
+
+func (r *Resource) HandleIntRequest(ctx context.Context, foo, bar string, f func(ls ctxlock.State) int) int {
+	// Same, but returns an int instead of a string.
+	// It must not allocate with the checked implementation.
+	mu := ctxlock.Lock(ctx, &r.mu)
+	defer mu.Unlock()
+	r.foo = foo
+	r.bar = bar
+	return f(mu.State())
+}
+
+func trimPanicMessage(r any) string {
+	msg := fmt.Sprintf("%v", r)
+	msg = strings.TrimSpace(msg)
+	if i := strings.IndexByte(msg, '\n'); i >= 0 {
+		return msg[:i]
+	}
+	return msg
+}