util/ctxlock: rename ctxlock.Context to ctxlock.State

Also add additional tests to verify that the unchecked implementation
is allocation-free.

Updates #12614

Signed-off-by: Nick Khyl <nickk@tailscale.com>

util/ctxlock/ctx_checked.go

@@ -1,51 +0,0 @@
-// Copyright (c) Tailscale Inc & AUTHORS
-// SPDX-License-Identifier: BSD-3-Clause
-
-// This file exports default, unoptimized implementation of the [Context] that includes runtime checks.
-// It is used unless the build tag ts_omit_ctxlock_checks is set.
-
-//go:build !ts_omit_ctxlock_checks
-
-package ctxlock
-
-import (
-	"context"
-	"sync"
-)
-
-// Context is a [context.Context] that can carry a [sync.Mutex] lock state.
-// Calling [Context.Unlock] on a [Context] unlocks the mutex locked by the context, if any.
-// It is a runtime error to call [Context.Unlock] more than once,
-// or use a [Context] after calling [Context.Unlock].
-type Context struct {
-	*checked
-}
-
-// None returns a [Context] that carries no mutex lock state and an empty [context.Context].
-//
-// It is typically used by top-level callers that do not have a parent context to pass in,
-// and is a shorthand for [Context]([context.Background]).
-func None() Context {
-	return Context{}
-}
-
-// Wrap returns a derived [Context] that wraps the provided [context.Context].
-//
-// It is typically used by callers that already have a [context.Context],
-// which may or may not be a [Context] tracking a mutex lock state.
-func Wrap(parent context.Context) Context {
-	return Context{wrapChecked(parent)}
-}
-
-// Lock returns a derived [Context] that wraps the provided [context.Context]
-// and carries the mutex lock state.
-//
-// It locks the mutex unless it is already held by the parent or an ancestor [Context].
-// It is a runtime error to pass a nil mutex or to unlock the parent context
-// before the returned one.
-func Lock[T context.Context](parent T, mu *sync.Mutex) Context {
-	if parent, ok := any(parent).(Context); ok {
-		return Context{lockChecked(parent.checked, mu)}
-	}
-	return Context{lockChecked(wrapChecked(parent), mu)}
-}

util/ctxlock/ctx_unchecked.go

@@ -1,33 +0,0 @@
-// Copyright (c) Tailscale Inc & AUTHORS
-// SPDX-License-Identifier: BSD-3-Clause
-
-// This file exports optimized implementation of the [Context] that omits runtime checks.
-// It is used when the build tag ts_omit_ctxlock_checks is set.
-
-//go:build ts_omit_ctxlock_checks
-
-package ctxlock
-
-import (
-	"context"
-	"sync"
-)
-
-type Context struct {
-	unchecked
-}
-
-func None() Context {
-	return Context{}
-}
-
-func Wrap(parent context.Context) Context {
-	return Context{wrapUnchecked(parent)}
-}
-
-func Lock[T context.Context](parent T, mu *sync.Mutex) Context {
-	if parent, ok := any(parent).(Context); ok {
-		return Context{lockUnchecked(parent.unchecked, mu)}
-	}
-	return Context{lockUnchecked(wrapUnchecked(parent), mu)}
-}

util/ctxlock/doc_test.go

@@ -7,6 +7,7 @@ import (
 	"context"
 	"fmt"
 	"sync"
+	"testing"
 
 	"tailscale.com/util/ctxlock"
 )
@@ -16,50 +17,63 @@ type Resource struct {
 	foo, bar string
 }
 
-func (r *Resource) GetFoo(ctx ctxlock.Context) string {
+func (r *Resource) GetFoo(ctx ctxlock.State) string {
-	defer ctxlock.Lock(ctx, &r.mu).Unlock() // Lock the mutex if not already held.
+	// Lock the mutex if not already held.
+	defer ctxlock.Lock(ctx, &r.mu).Unlock()
 	return r.foo
 }
 
-func (r *Resource) SetFoo(ctx ctxlock.Context, foo string) {
+func (r *Resource) SetFoo(ctx ctxlock.State, foo string) {
-	defer ctxlock.Lock(ctx, &r.mu).Unlock()
+	// 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
 }
 
-func (r *Resource) GetBar(ctx ctxlock.Context) string {
+func (r *Resource) GetBar(ctx ctxlock.State) string {
-	ctx = ctxlock.Lock(ctx, &r.mu)
+	defer ctxlock.Lock(ctx, &r.mu).Unlock()
-	defer ctx.Unlock() // If you prefer it this way.
 	return r.bar
 }
 
-func (r *Resource) SetBar(ctx ctxlock.Context, bar string) {
+func (r *Resource) SetBar(ctx ctxlock.State, bar string) {
 	defer ctxlock.Lock(ctx, &r.mu).Unlock()
 	r.bar = bar
 }
 
-func (r *Resource) WithLock(ctx ctxlock.Context, f func(ctx ctxlock.Context)) {
+func (r *Resource) WithLock(ctx ctxlock.State, f func(ctx ctxlock.State)) {
-	// Lock the mutex if not already held, and get a new context.
+	// 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 context.
+	f(ctx) // Call the callback with the new lock state.
 }
 
-func (r *Resource) HandleRequest(ctx context.Context, foo, bar string, f func(ctx ctxlock.Context) string) string {
+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.Context].
+	// Same, but with a standard [context.Context] instead of [ctxlock.State].
 	// [ctxlock.Lock] is generic and works with both without allocating.
-	// The provided context can be used for cancellation, etc.
+	// The ctx can be used for cancellation, etc.
-	muCtx := ctxlock.Lock(ctx, &r.mu)
+	mu := ctxlock.Lock(ctx, &r.mu)
-	defer muCtx.Unlock()
+	defer mu.Unlock()
 	r.foo = foo
 	r.bar = bar
-	return f(muCtx)
+	return f(mu)
 }
 
-func ExampleContext() {
+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 ExampleState() {
 	var r Resource
 	r.SetFoo(ctxlock.None(), "foo")
 	r.SetBar(ctxlock.None(), "bar")
-	r.WithLock(ctxlock.None(), func(ctx ctxlock.Context) {
+	r.WithLock(ctxlock.None(), func(ctx ctxlock.State) {
 		// This callback is invoked with r's lock held,
 		// and ctx carries the lock state. This means we can safely call
 		// other methods on r using ctx without causing a deadlock.
@@ -69,11 +83,11 @@ func ExampleContext() {
 	// Output: foobar
 }
 
-func ExampleContext_twoResources() {
+func ExampleState_twoResources() {
 	var r1, r2 Resource
 	r1.SetFoo(ctxlock.None(), "foo")
 	r2.SetBar(ctxlock.None(), "bar")
-	r1.WithLock(ctxlock.None(), func(ctx ctxlock.Context) {
+	r1.WithLock(ctxlock.None(), func(ctx ctxlock.State) {
 		// Here, r1's lock is held, but r2's lock is not.
 		// So r2 will be locked when we call r2.GetBar(ctx).
 		r1.SetFoo(ctx, r1.GetFoo(ctx)+r2.GetBar(ctx))
@@ -82,26 +96,35 @@ func ExampleContext_twoResources() {
 	// Output: foobar
 }
 
-func ExampleContext_zeroValue() {
+func ExampleState_stdContext() {
-	var r1, r2 Resource
-	r1.SetFoo(ctxlock.Context{}, "foo")
-	r2.SetBar(ctxlock.Context{}, "bar")
-	r1.WithLock(ctxlock.Context{}, func(ctx ctxlock.Context) {
-		// Here, r1's lock is held, but r2's lock is not.
-		// So r2 will be locked when we call r2.GetBar(ctx).
-		r1.SetFoo(ctx, r1.GetFoo(ctx)+r2.GetBar(ctx))
-	})
-	fmt.Println(r1.GetFoo(ctxlock.Context{}))
-	// Output: foobar
-}
-
-func ExampleContext_stdContext() {
 	var r Resource
 	ctx := context.Background()
-	result := r.HandleRequest(ctx, "foo", "bar", func(ctx ctxlock.Context) string {
+	result := r.HandleRequest(ctx, "foo", "bar", func(ctx ctxlock.State) string {
 		// The r's lock 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)")
+	}
+
+	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.
+			return len(r.GetFoo(ctx) + r.GetBar(ctx))
+		})
+		if res != 6 {
+			t.Errorf("expected 6, got %d", res)
+		}
+	}); allocs != 0 {
+		t.Errorf("expected 0 allocs, got %f", allocs)
+	}
+}

util/ctxlock/state.go

@@ -11,6 +11,8 @@
 // It defaults to the checked implementation unless the ts_omit_ctxlock_checks build tag is set.
 package ctxlock
 
+// This file contains both the [checked] and [unchecked] implementations of [State].
+
 import (
 	"context"
 	"fmt"
@@ -25,29 +27,40 @@ func ctxKeyOf(mu *sync.Mutex) ctxKey {
 	return ctxKey{mu}
 }
 
-// checked is an implementation of [Context] that performs runtime checks
+// checked is an implementation of [State] that performs runtime checks
-// to ensure that the context is used correctly.
+// to ensure the correct order of locking and unlocking.
 //
-// Its zero value and a nil pointer carry no lock state and an empty [context.Context].
+// Its zero value and a nil pointer are valid and carry no lock state
+// and an empty [context.Context].
 type checked struct {
-	context.Context             // nil if the context does not carry a [context.Context]
+	context.Context // nil means an empty context
-	mu              *sync.Mutex // nil if the context does not carry a mutex lock state
+
-	parent          *checked    // nil if the context owns the lock
+	// mu is the mutex tracked by this state,
-	unlocked        bool        // whether [checked.Unlock] was called
+	// or nil if it wasn't created with [Lock].
+	mu *sync.Mutex
+
+	// parent is an ancestor State associated with the same mutex.
+	// It may or may not own the lock (the lock could be held by a further ancestor).
+	// The parent is nil if this State is the root of the hierarchy,
+	// meaning it owns the lock.
+	parent *checked
+
+	// unlocked is whether [checked.Unlock] was called on this state.
+	unlocked bool
 }
 
-func wrapChecked(parent context.Context) *checked {
+func fromContextChecked(ctx context.Context) *checked {
-	return &checked{parent, nil, nil, false}
+	return &checked{ctx, nil, nil, false}
 }
 
 func lockChecked(parent *checked, mu *sync.Mutex) *checked {
 	panicIfNil(mu)
-	if parentLockCtx, ok := parent.Value(ctxKeyOf(mu)).(*checked); ok {
+	if parentState, ok := parent.Value(ctxKeyOf(mu)).(*checked); ok {
 		if appearsUnlocked(mu) {
-			// The parent still owns the lock, but the mutex is unlocked.
+			// The parent is already unlocked, but the mutex is not.
-			panic("mu is spuriously unlocked")
+			panic(fmt.Sprintf("%T is spuriously unlocked", mu))
 		}
-		return &checked{parent, mu, parentLockCtx, false}
+		return &checked{parent, mu, parentState, false}
 	}
 	mu.Lock()
 	return &checked{parent, mu, nil, false}
@@ -80,48 +93,55 @@ func (c *checked) Err() error {
 func (c *checked) Value(key any) any {
 	c.panicIfUnlocked()
 	if c == nil {
+		// No-op; zero state.
 		return nil
 	}
-	if key == ctxKeyOf(c.mu) {
+	if key, ok := key.(ctxKey); ok && key.Mutex == c.mu {
+		// This is the mutex tracked by this state.
 		return c
 	}
-	return c.Context.Value(key)
+	if c.Context != nil {
+		// Forward the call to the parent context,
+		// which may or may not be a [checked] state.
+		return c.Context.Value(key)
+	}
+	return nil
 }
 
 func (c *checked) Unlock() {
 	switch {
 	case c == nil:
-		// No-op; zero context.
+		// No-op; zero state.
 		return
 	case c.unlocked:
 		panic("already unlocked")
 	case c.mu == nil:
-		// No-op; the context does not track a mutex lock state,
+		// No-op; the state does not track a mutex lock state,
-		// such as when it was created with [noneChecked] or [wrapChecked].
+		// meaning it was not created with [Lock].
 	case c.parent == nil:
-		// We own the lock; let's unlock it.
+		// The state own the mutex's lock; we must unlock it.
 		// This triggers a fatal error if the mutex is already unlocked.
 		c.mu.Unlock()
 	case c.parent.unlocked:
-		// The parent context is already unlocked.
+		// The parent state is already unlocked.
 		// The mutex may or may not be locked;
 		// something else may have already locked it.
 		panic("parent already unlocked")
 	case appearsUnlocked(c.mu):
 		// The mutex itself is unlocked,
-		// even though the parent context is still locked.
+		// even though the parent state is still locked.
-		// It may be unlocked by an ancestor context
+		// It may be unlocked by an ancestor state
 		// or by something else entirely.
 		panic("mutex is not locked")
 	default:
 		// No-op; a parent or ancestor will handle unlocking.
 	}
-	c.unlocked = true // mark this context as unlocked
+	c.unlocked = true // mark this state as unlocked
 }
 
 func (c *checked) panicIfUnlocked() {
 	if c != nil && c.unlocked {
-		panic("use of context after unlock")
+		panic("use after unlock")
 	}
 }
 
@@ -131,23 +151,26 @@ func panicIfNil[T comparable](v T) {
 	}
 }
 
-// unchecked is an implementation of [Context] that trades runtime checks for performance.
+// unchecked is an implementation of [State] that trades runtime checks for performance.
 //
 // Its zero value carries no mutex lock state and an empty [context.Context].
 type unchecked struct {
-	context.Context             // nil if the context does not carry a [context.Context]
+	context.Context             // nil means an empty context
-	mu              *sync.Mutex // non-nil if locked by this context
+	mu              *sync.Mutex // non-nil if owned by this state
 }
 
-func wrapUnchecked(parent context.Context) unchecked {
+func fromContextUnchecked(ctx context.Context) unchecked {
-	return unchecked{parent, nil}
+	return unchecked{ctx, nil}
 }
 
 func lockUnchecked(parent unchecked, mu *sync.Mutex) unchecked {
 	if parent.Value(ctxKeyOf(mu)) == nil {
+		// There's no ancestor state associated with this mutex,
+		// so we can lock it.
 		mu.Lock()
 	} else {
-		mu = nil // already locked by a parent/ancestor
+		// The mutex is already locked by a parent/ancestor state.
+		mu = nil
 	}
 	return unchecked{parent.Context, mu}
 }
@@ -174,7 +197,7 @@ func (c unchecked) Err() error {
 }
 
 func (c unchecked) Value(key any) any {
-	if key == ctxKeyOf(c.mu) {
+	if key, ok := key.(ctxKey); ok && key.Mutex == c.mu {
 		return key
 	}
 	if c.Context == nil {

util/ctxlock/state_checked.go

@@ -0,0 +1,53 @@
+// Copyright (c) Tailscale Inc & AUTHORS
+// SPDX-License-Identifier: BSD-3-Clause
+
+// This file exports default, unoptimized implementation of the [State] that includes runtime checks.
+// It is used unless the build tag ts_omit_ctxlock_checks is set.
+
+//go:build !ts_omit_ctxlock_checks
+
+package ctxlock
+
+import (
+	"context"
+	"sync"
+)
+
+// Checked indicates whether runtime checks are enabled for this package.
+const Checked = true
+
+// State carries the lock state of zero or more mutexes and an optional [context.Context].
+// Its zero value is valid and represents an unlocked state and an empty context.
+//
+// Calling [Lock] returns a derived State with the specified mutex locked. The State is considered
+// the owner of the lock if it wasn't already acquired by a parent State. Calling [State.Unlock]
+// releases the lock owned by the state. It is a runtime error to call Unlock more than once,
+// to use the State after it has been unlocked, or to unlock a parent State before its child.
+type State struct {
+	*checked
+}
+
+// None returns a [State] that carries no lock state and an empty [context.Context].
+func None() State {
+	return State{}
+}
+
+// FromContext returns a [State] that carries the same lock state as the provided [context.Context].
+//
+// It is typically used by methods that already accept a [context.Context] for cancellation or deadline
+// management, and would like to use it for locking as well.
+func FromContext(ctx context.Context) State {
+	return State{fromContextChecked(ctx)}
+}
+
+// Lock acquires the specified mutex and becomes its owner, unless it is already held by a parent.
+// The parent can be either a [State] or a [context.Context]. A zero [State] is a valid parent.
+// It returns a new [State] that augments the parent with the additional lock state.
+//
+// It is a runtime error to pass a nil mutex or to unlock the parent state before the returned one.
+func Lock[T context.Context](parent T, mu *sync.Mutex) State {
+	if parent, ok := any(parent).(State); ok {
+		return State{lockChecked(parent.checked, mu)}
+	}
+	return State{lockChecked(fromContextChecked(parent), mu)}
+}

util/ctxlock/state_test.go

@@ -11,35 +11,72 @@ import (
 	"tailscale.com/util/ctxkey"
 )
 
-type ctx interface {
+type state interface {
 	context.Context
 	Unlock()
 }
 
-type impl[T ctx] struct {
+type impl[T state] struct {
-	None func() T
+	None        func() T
-	Wrap func(context.Context) T
+	FromContext func(context.Context) T
-	Lock func(T, *sync.Mutex) T
+	Lock        func(T, *sync.Mutex) T
+	LockCtx     func(context.Context, *sync.Mutex) T
 }
 
 var (
-	exportedImpl = impl[Context]{
+	exportedImpl = impl[State]{
-		None: None,
+		None:        None,
-		Wrap: Wrap,
+		FromContext: FromContext,
-		Lock: Lock[Context],
+		Lock:        Lock[State],
+		LockCtx:     Lock[context.Context],
 	}
 	checkedImpl = impl[*checked]{
-		None: func() *checked { return nil },
+		None:        func() *checked { return nil },
-		Wrap: wrapChecked,
+		FromContext: fromContextChecked,
-		Lock: lockChecked,
+		Lock:        lockChecked,
+		LockCtx: func(ctx context.Context, mu *sync.Mutex) *checked {
+			return lockChecked(fromContextChecked(ctx), mu)
+		},
 	}
 	uncheckedImpl = impl[unchecked]{
-		None: func() unchecked { return unchecked{} },
+		None:        func() unchecked { return unchecked{} },
-		Wrap: wrapUnchecked,
+		FromContext: fromContextUnchecked,
-		Lock: lockUnchecked,
+		Lock:        lockUnchecked,
+		LockCtx: func(ctx context.Context, mu *sync.Mutex) unchecked {
+			return lockUnchecked(fromContextUnchecked(ctx), mu)
+		},
 	}
 )
 
+// BenchmarkLockUnlock benchmarks the performance of locking and unlocking a mutex.
+func BenchmarkLockUnlock(b *testing.B) {
+	var mu sync.Mutex
+	b.Run("Exported", func(b *testing.B) {
+		benchmarkLockUnlock(b, exportedImpl)
+	})
+	b.Run("Checked", func(b *testing.B) {
+		benchmarkLockUnlock(b, checkedImpl)
+	})
+	b.Run("Unchecked", func(b *testing.B) {
+		benchmarkLockUnlock(b, uncheckedImpl)
+	})
+	b.Run("Reference", func(b *testing.B) {
+		for b.Loop() {
+			mu.Lock()
+			mu.Unlock()
+		}
+	})
+}
+
+func benchmarkLockUnlock[T state](b *testing.B, impl impl[T]) {
+	var mu sync.Mutex
+	for b.Loop() {
+		ctx := impl.Lock(impl.None(), &mu)
+		ctx.Unlock()
+	}
+}
+
+// BenchmarkReentrance benchmarks the performance of reentrant locking and unlocking.
 func BenchmarkReentrance(b *testing.B) {
 	var mu sync.Mutex
 
@@ -65,7 +102,7 @@ func BenchmarkReentrance(b *testing.B) {
 	})
 }
 
-func benchmarkReentrance[T ctx](b *testing.B, impl impl[T]) {
+func benchmarkReentrance[T state](b *testing.B, impl impl[T]) {
 	var mu sync.Mutex
 	for b.Loop() {
 		parent := impl.Lock(impl.None(), &mu)
@@ -77,14 +114,16 @@ func benchmarkReentrance[T ctx](b *testing.B, impl impl[T]) {
 	}
 }
 
+// BenchmarkGenericLock benchmarks the performance of the generic [Lock] function
+// that works with both [State] and [context.Context].
 func BenchmarkGenericLock(b *testing.B) {
 	// Does not allocate with --tags=ts_omit_ctxlock_checks.
-	b.Run("ZeroContext", func(b *testing.B) {
+	b.Run("State", func(b *testing.B) {
 		var mu sync.Mutex
-		var ctx Context
+		var ctx State
 		for b.Loop() {
 			parent := Lock(ctx, &mu)
-			func(ctx Context) {
+			func(ctx State) {
 				child := Lock(ctx, &mu)
 				child.Unlock()
 			}(parent)
@@ -96,7 +135,7 @@ func BenchmarkGenericLock(b *testing.B) {
 		ctx := context.Background()
 		for b.Loop() {
 			parent := Lock(ctx, &mu)
-			func(ctx Context) {
+			func(ctx State) {
 				child := Lock(ctx, &mu)
 				child.Unlock()
 			}(parent)
@@ -105,6 +144,55 @@ func BenchmarkGenericLock(b *testing.B) {
 	})
 }
 
+// TestUncheckedAllocFree tests that the exported implementation of [State] does not allocate memory
+// when the ts_omit_ctxlock_checks build tag is set.
+func TestUncheckedAllocFree(t *testing.T) {
+	if Checked {
+		t.Skip("Exported implementation is not alloc-free (use --tags=ts_omit_ctxlock_checks)")
+	}
+	t.Run("Simple/WithState", func(t *testing.T) {
+		var mu sync.Mutex
+		mustNotAllocate(t, func() {
+			ctx := Lock(None(), &mu)
+			ctx.Unlock()
+		})
+	})
+
+	t.Run("Simple/WithContext", func(t *testing.T) {
+		var mu sync.Mutex
+		ctx := context.Background()
+		mustNotAllocate(t, func() {
+			ctx := Lock(ctx, &mu)
+			ctx.Unlock()
+		})
+	})
+
+	t.Run("Reentrant/WithState", func(t *testing.T) {
+		var mu sync.Mutex
+		mustNotAllocate(t, func() {
+			parent := Lock(None(), &mu)
+			func(ctx State) {
+				child := Lock(parent, &mu)
+				child.Unlock()
+			}(parent)
+			parent.Unlock()
+		})
+	})
+
+	t.Run("Reentrant/WithContext", func(t *testing.T) {
+		var mu sync.Mutex
+		ctx := context.Background()
+		mustNotAllocate(t, func() {
+			parent := Lock(ctx, &mu)
+			func(ctx State) {
+				child := Lock(ctx, &mu)
+				child.Unlock()
+			}(parent)
+			parent.Unlock()
+		})
+	})
+}
+
 func TestHappyPath(t *testing.T) {
 	t.Run("Exported", func(t *testing.T) {
 		testHappyPath(t, exportedImpl)
@@ -119,7 +207,7 @@ func TestHappyPath(t *testing.T) {
 	})
 }
 
-func testHappyPath[T ctx](t *testing.T, impl impl[T]) {
+func testHappyPath[T state](t *testing.T, impl impl[T]) {
 	var mu sync.Mutex
 	parent := impl.Lock(impl.None(), &mu)
 	wantLocked(t, &mu) // mu is locked by parent
@@ -128,9 +216,9 @@ func testHappyPath[T ctx](t *testing.T, impl impl[T]) {
 	wantLocked(t, &mu) // mu is still locked by parent
 
 	var mu2 sync.Mutex
-	context2 := impl.Lock(child, &mu2)
+	ls2 := impl.Lock(child, &mu2)
-	wantLocked(t, &mu2)   // mu2 is locked by context2
+	wantLocked(t, &mu2)   // mu2 is locked by ls2
-	context2.Unlock()     // unlocks mu2
+	ls2.Unlock()          // unlocks mu2
 	wantUnlocked(t, &mu2) // mu2 is now unlocked
 
 	child.Unlock()     // noop
@@ -140,46 +228,50 @@ func testHappyPath[T ctx](t *testing.T, impl impl[T]) {
 	wantUnlocked(t, &mu) // mu is now unlocked
 }
 
-func TestWrappedLockContext(t *testing.T) {
+func TestContextWrapping(t *testing.T) {
 	t.Run("Exported", func(t *testing.T) {
-		testWrappedLockContext(t, exportedImpl)
+		testContextWrapping(t, exportedImpl)
 	})
 
 	t.Run("Checked", func(t *testing.T) {
-		testWrappedLockContext(t, checkedImpl)
+		testContextWrapping(t, checkedImpl)
 	})
 
 	t.Run("Unchecked", func(t *testing.T) {
-		testWrappedLockContext(t, uncheckedImpl)
+		testContextWrapping(t, uncheckedImpl)
 	})
 }
 
-func testWrappedLockContext[T ctx](t *testing.T, impl impl[T]) {
+func testContextWrapping[T state](t *testing.T, impl impl[T]) {
+	// Create a [context.Context] with a value set in it.
 	wantValue := "value"
 	key := ctxkey.New("key", "")
 	ctxWithValue := key.WithValue(context.Background(), wantValue)
-	root := impl.Wrap(ctxWithValue)
 
 	var mu sync.Mutex
-	parent := impl.Lock(root, &mu)
+	parent := impl.LockCtx(ctxWithValue, &mu)
 	wantLocked(t, &mu) // mu is locked by parent
 
-	// Wrap the parent context as if it were a regular [context.Context],
+	// Let's assume that we want to call a function that takes a [context.Context].
-	// then create a child context from it.
+	// [State] is a valid [context.Context], so we can pass it to the function.
-	// The child should still recognize the parent as the mutex owner,
+	ctx := context.Context(parent)
-	// and not panic or deadlock attempting to lock it again.
+	// If / when necessary, we can convert it back to a [State].
-	wrapped := impl.Wrap(parent)
+	// The [State] should carry the same lock state as the parent context.
-	child := impl.Lock(wrapped, &mu)
+	parentDup := impl.FromContext(ctx)
 
-	// We should be able to access the value set in the root context.
+	// We can then create and use a child [State].
+	child := impl.Lock(parentDup, &mu)
+
+	// It still carries all the original context values...
 	if gotValue := key.Value(child); gotValue != wantValue {
 		t.Errorf("key.Value() = %s; want %s", gotValue, wantValue)
 	}
 
+	// ... and the lock state.
 	child.Unlock()     // no-op; mu is owned by parent
 	wantLocked(t, &mu) // mu is still locked by parent
 
-	wrapped.Unlock()   // no-op; mu is owned by parent
+	parentDup.Unlock() // no-op; mu is owned by parent
 	wantLocked(t, &mu) // mu is still locked by parent
 
 	parent.Unlock()      // unlocks mu
@@ -198,7 +290,7 @@ func TestUseUnlockedParent_Checked(t *testing.T) {
 	parent := impl.Lock(impl.None(), &mu)
 	parent.Unlock()      // unlocks mu
 	wantUnlocked(t, &mu) // mu is now unlocked
-	wantPanic(t, "use of context after unlock", func() { impl.Lock(parent, &mu) })
+	wantPanic(t, "use after unlock", func() { impl.Lock(parent, &mu) })
 }
 
 func TestUseUnlockedMutex_Checked(t *testing.T) {
@@ -207,14 +299,14 @@ func TestUseUnlockedMutex_Checked(t *testing.T) {
 	var mu sync.Mutex
 	parent := impl.Lock(impl.None(), &mu)
 	mu.Unlock() // unlock mu directly without unlocking parent
-	wantPanic(t, "mu is spuriously unlocked", func() { impl.Lock(parent, &mu) })
+	wantPanic(t, "*sync.Mutex is spuriously unlocked", func() { impl.Lock(parent, &mu) })
 }
 
 func TestUnlockParentFirst_Checked(t *testing.T) {
 	impl := checkedImpl
 
 	var mu sync.Mutex
-	parent := impl.Lock(impl.Wrap(context.Background()), &mu)
+	parent := impl.Lock(impl.FromContext(context.Background()), &mu)
 	child := impl.Lock(parent, &mu)
 
 	parent.Unlock()      // unlocks mu
@@ -231,7 +323,7 @@ func TestUnlockTwice_Checked(t *testing.T) {
 	}
 
 	t.Run("Wrapped", func(t *testing.T) {
-		unlockTwice(t, impl.Wrap(context.Background()))
+		unlockTwice(t, impl.FromContext(context.Background()))
 	})
 	t.Run("Locked", func(t *testing.T) {
 		var mu sync.Mutex
@@ -242,7 +334,7 @@ func TestUnlockTwice_Checked(t *testing.T) {
 		var mu sync.Mutex
 		ctx := impl.Lock(impl.None(), &mu)
 		ctx.Unlock() // unlocks mu
-		mu.Lock()    // re-locks mu, but not by the context
+		mu.Lock()    // re-locks mu, but not by the state
 		wantPanic(t, "already unlocked", ctx.Unlock)
 	})
 	t.Run("Child", func(t *testing.T) {
@@ -257,14 +349,14 @@ func TestUnlockTwice_Checked(t *testing.T) {
 		parent := impl.Lock(impl.None(), &mu)
 		child := impl.Lock(parent, &mu)
 		parent.Unlock()
-		mu.Lock() // re-locks mu, but not the parent context
+		mu.Lock() // re-locks mu, but not the parent state
 		wantPanic(t, "parent already unlocked", child.Unlock)
 	})
 	t.Run("Child/WithManualUnlock", func(t *testing.T) {
 		var mu sync.Mutex
 		parent := impl.Lock(impl.None(), &mu)
 		child := impl.Lock(parent, &mu)
-		mu.Unlock() // unlocks mu, but not the parent context
+		mu.Unlock() // unlocks mu, but not the parent state
 		wantPanic(t, "mutex is not locked", child.Unlock)
 	})
 	t.Run("Grandchild", func(t *testing.T) {
@@ -282,50 +374,50 @@ func TestUseUnlocked_Checked(t *testing.T) {
 	impl := checkedImpl
 
 	var mu sync.Mutex
-	ctx := lockChecked(impl.None(), &mu)
+	state := lockChecked(impl.None(), &mu)
-	ctx.Unlock()
+	state.Unlock()
 
-	// All of these should panic since the context is already unlocked.
+	// All of these should panic since the state is already unlocked.
-	wantPanic(t, "", func() { ctx.Deadline() })
+	wantPanic(t, "", func() { state.Deadline() })
-	wantPanic(t, "", func() { ctx.Done() })
+	wantPanic(t, "", func() { state.Done() })
-	wantPanic(t, "", func() { ctx.Err() })
+	wantPanic(t, "", func() { state.Err() })
-	wantPanic(t, "", func() { ctx.Unlock() })
+	wantPanic(t, "", func() { state.Unlock() })
-	wantPanic(t, "", func() { ctx.Value("key") })
+	wantPanic(t, "", func() { state.Value("key") })
 }
 
-func TestUseNoneContext(t *testing.T) {
+func TestUseZeroState(t *testing.T) {
 	t.Run("Exported", func(t *testing.T) {
-		testUseEmptyContext(t, exportedImpl.None, exportedImpl)
+		testUseEmptyState(t, exportedImpl.None, exportedImpl)
 	})
 	t.Run("Checked", func(t *testing.T) {
-		testUseEmptyContext(t, checkedImpl.None, checkedImpl)
+		testUseEmptyState(t, checkedImpl.None, checkedImpl)
 	})
 	t.Run("Unchecked", func(t *testing.T) {
-		testUseEmptyContext(t, uncheckedImpl.None, uncheckedImpl)
+		testUseEmptyState(t, uncheckedImpl.None, uncheckedImpl)
 	})
 }
 
 func TestUseWrappedBackground(t *testing.T) {
 	t.Run("Exported", func(t *testing.T) {
-		testUseEmptyContext(t, getWrappedBackground(t, exportedImpl), exportedImpl)
+		testUseEmptyState(t, getWrappedBackground(t, exportedImpl), exportedImpl)
 	})
 	t.Run("Checked", func(t *testing.T) {
-		testUseEmptyContext(t, getWrappedBackground(t, checkedImpl), checkedImpl)
+		testUseEmptyState(t, getWrappedBackground(t, checkedImpl), checkedImpl)
 	})
 	t.Run("Unchecked", func(t *testing.T) {
-		testUseEmptyContext(t, getWrappedBackground(t, uncheckedImpl), uncheckedImpl)
+		testUseEmptyState(t, getWrappedBackground(t, uncheckedImpl), uncheckedImpl)
 	})
 }
 
-func getWrappedBackground[T ctx](t *testing.T, impl impl[T]) func() T {
+func getWrappedBackground[T state](t *testing.T, impl impl[T]) func() T {
 	t.Helper()
 	return func() T {
-		return impl.Wrap(context.Background())
+		return impl.FromContext(context.Background())
 	}
 }
 
-func testUseEmptyContext[T ctx](t *testing.T, getCtx func() T, impl impl[T]) {
+func testUseEmptyState[T state](t *testing.T, getCtx func() T, impl impl[T]) {
-	// Using a None context must not panic or deadlock.
+	// Using aan empty [State] must not panic or deadlock.
 	// It should also behave like [context.Background].
 	for range 2 {
 		ctx := getCtx()
@@ -372,3 +464,11 @@ func wantUnlocked(t *testing.T, m *sync.Mutex) {
 	}
 	m.Unlock()
 }
+
+func mustNotAllocate(t *testing.T, steps func()) {
+	t.Helper()
+	const runs = 1000
+	if allocs := testing.AllocsPerRun(runs, steps); allocs != 0 {
+		t.Errorf("expected 0 allocs, got %f", allocs)
+	}
+}

util/ctxlock/state_unchecked.go

@@ -0,0 +1,35 @@
+// Copyright (c) Tailscale Inc & AUTHORS
+// SPDX-License-Identifier: BSD-3-Clause
+
+// This file exports optimized implementation of the [State] that omits runtime checks.
+// It is used when the build tag ts_omit_ctxlock_checks is set.
+
+//go:build ts_omit_ctxlock_checks
+
+package ctxlock
+
+import (
+	"context"
+	"sync"
+)
+
+const Checked = false
+
+type State struct {
+	unchecked
+}
+
+func None() State {
+	return State{}
+}
+
+func FromContext(parent context.Context) State {
+	return State{fromContextUnchecked(parent)}
+}
+
+func Lock[T context.Context](parent T, mu *sync.Mutex) State {
+	if parent, ok := any(parent).(State); ok {
+		return State{lockUnchecked(parent.unchecked, mu)}
+	}
+	return State{lockUnchecked(fromContextUnchecked(parent), mu)}
+}