mirror of
https://github.com/tailscale/tailscale.git
synced 2025-06-11 10:18:35 +00:00
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>
This commit is contained in:
Nick Khyl
parent
b802e8674c
commit
64e5da8024
@ -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)}
|
||||
}
|
||||
@ -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)}
|
||||
}
|
||||
@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
@ -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 {
|
||||
53
util/ctxlock/state_checked.go
Normal file
53
util/ctxlock/state_checked.go
Normal file
@ -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)}
|
||||
}
|
||||
@ -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)
|
||||
}
|
||||
}
|
||||
35
util/ctxlock/state_unchecked.go
Normal file
35
util/ctxlock/state_unchecked.go
Normal file
@ -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)}
|
||||
}
|
||||
Loading…
x
Reference in New Issue
Block a user