diff --git a/util/ctxlock/ctx_checked.go b/util/ctxlock/ctx_checked.go deleted file mode 100644 index a214294bb..000000000 --- a/util/ctxlock/ctx_checked.go +++ /dev/null @@ -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)} -} diff --git a/util/ctxlock/ctx_unchecked.go b/util/ctxlock/ctx_unchecked.go deleted file mode 100644 index faa662cd1..000000000 --- a/util/ctxlock/ctx_unchecked.go +++ /dev/null @@ -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)} -} diff --git a/util/ctxlock/doc_test.go b/util/ctxlock/doc_test.go index 0bc5592c2..a6b0de407 100644 --- a/util/ctxlock/doc_test.go +++ b/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 { - defer ctxlock.Lock(ctx, &r.mu).Unlock() // Lock the mutex if not already held. +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 } -func (r *Resource) SetFoo(ctx ctxlock.Context, foo string) { - defer ctxlock.Lock(ctx, &r.mu).Unlock() +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 } -func (r *Resource) GetBar(ctx ctxlock.Context) string { - ctx = ctxlock.Lock(ctx, &r.mu) - defer ctx.Unlock() // If you prefer it this way. +func (r *Resource) GetBar(ctx ctxlock.State) string { + defer ctxlock.Lock(ctx, &r.mu).Unlock() 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)) { - // Lock the mutex if not already held, and get a new context. +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 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 { - // Same, but with a standard [context.Context] instead of [ctxlock.Context]. +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 provided context can be used for cancellation, etc. - muCtx := ctxlock.Lock(ctx, &r.mu) - defer muCtx.Unlock() + // 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(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() { - 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() { +func ExampleState_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) + } +} diff --git a/util/ctxlock/ctx.go b/util/ctxlock/state.go similarity index 60% rename from util/ctxlock/ctx.go rename to util/ctxlock/state.go index ae82ed0b2..4ec9857e6 100644 --- a/util/ctxlock/ctx.go +++ b/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 -// to ensure that the context is used correctly. +// checked is an implementation of [State] that performs runtime checks +// 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] - mu *sync.Mutex // nil if the context does not carry a mutex lock state - parent *checked // nil if the context owns the lock - unlocked bool // whether [checked.Unlock] was called + context.Context // nil means an empty context + + // mu is the mutex tracked by this state, + // 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 { - return &checked{parent, nil, nil, false} +func fromContextChecked(ctx context.Context) *checked { + 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. - panic("mu is spuriously unlocked") + // The parent is already unlocked, but the mutex is not. + 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, - // such as when it was created with [noneChecked] or [wrapChecked]. + // No-op; the state does not track a mutex lock state, + // 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. - // It may be unlocked by an ancestor context + // even though the parent state is still locked. + // 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] - mu *sync.Mutex // non-nil if locked by this context + context.Context // nil means an empty context + mu *sync.Mutex // non-nil if owned by this state } -func wrapUnchecked(parent context.Context) unchecked { - return unchecked{parent, nil} +func fromContextUnchecked(ctx context.Context) unchecked { + 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 { diff --git a/util/ctxlock/state_checked.go b/util/ctxlock/state_checked.go new file mode 100644 index 000000000..4705b9c19 --- /dev/null +++ b/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)} +} diff --git a/util/ctxlock/ctx_test.go b/util/ctxlock/state_test.go similarity index 53% rename from util/ctxlock/ctx_test.go rename to util/ctxlock/state_test.go index 5d4a77066..e9b747c97 100644 --- a/util/ctxlock/ctx_test.go +++ b/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 { - None func() T - Wrap func(context.Context) T - Lock func(T, *sync.Mutex) T +type impl[T state] struct { + None func() T + FromContext func(context.Context) T + Lock func(T, *sync.Mutex) T + LockCtx func(context.Context, *sync.Mutex) T } var ( - exportedImpl = impl[Context]{ - None: None, - Wrap: Wrap, - Lock: Lock[Context], + exportedImpl = impl[State]{ + None: None, + FromContext: FromContext, + Lock: Lock[State], + LockCtx: Lock[context.Context], } checkedImpl = impl[*checked]{ - None: func() *checked { return nil }, - Wrap: wrapChecked, - Lock: lockChecked, + None: func() *checked { return nil }, + FromContext: fromContextChecked, + Lock: lockChecked, + LockCtx: func(ctx context.Context, mu *sync.Mutex) *checked { + return lockChecked(fromContextChecked(ctx), mu) + }, } uncheckedImpl = impl[unchecked]{ - None: func() unchecked { return unchecked{} }, - Wrap: wrapUnchecked, - Lock: lockUnchecked, + None: func() unchecked { return unchecked{} }, + FromContext: fromContextUnchecked, + 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) - wantLocked(t, &mu2) // mu2 is locked by context2 - context2.Unlock() // unlocks mu2 + ls2 := impl.Lock(child, &mu2) + wantLocked(t, &mu2) // mu2 is locked by ls2 + 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], - // then create a child context from it. - // The child should still recognize the parent as the mutex owner, - // and not panic or deadlock attempting to lock it again. - wrapped := impl.Wrap(parent) - child := impl.Lock(wrapped, &mu) + // Let's assume that we want to call a function that takes a [context.Context]. + // [State] is a valid [context.Context], so we can pass it to the function. + ctx := context.Context(parent) + // If / when necessary, we can convert it back to a [State]. + // The [State] should carry the same lock state as the parent context. + 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) - ctx.Unlock() + state := lockChecked(impl.None(), &mu) + state.Unlock() - // All of these should panic since the context is already unlocked. - wantPanic(t, "", func() { ctx.Deadline() }) - wantPanic(t, "", func() { ctx.Done() }) - wantPanic(t, "", func() { ctx.Err() }) - wantPanic(t, "", func() { ctx.Unlock() }) - wantPanic(t, "", func() { ctx.Value("key") }) + // All of these should panic since the state is already unlocked. + wantPanic(t, "", func() { state.Deadline() }) + wantPanic(t, "", func() { state.Done() }) + wantPanic(t, "", func() { state.Err() }) + wantPanic(t, "", func() { state.Unlock() }) + 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]) { - // Using a None context must not panic or deadlock. +func testUseEmptyState[T state](t *testing.T, getCtx func() T, impl impl[T]) { + // 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) + } +} diff --git a/util/ctxlock/state_unchecked.go b/util/ctxlock/state_unchecked.go new file mode 100644 index 000000000..c55150b75 --- /dev/null +++ b/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)} +}