util/ctxlock: make zero Context a valid, empty context

Updates #12614

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

util/ctxlock/ctx.go

@@ -14,12 +14,9 @@ package ctxlock
 import (
 	"context"
 	"fmt"
+	"reflect"
 	"sync"
-)
-
-var (
-	noneCtx       = context.Background()
-	noneUnchecked = unchecked{noneCtx, nil}
+	"time"
 )
 
 type ctxKey struct{ *sync.Mutex }
@@ -30,36 +27,60 @@ func ctxKeyOf(mu *sync.Mutex) ctxKey {
 
 // checked is an implementation of [Context] that performs runtime checks
 // to ensure that the context is used correctly.
+//
+// Its zero value and a nil pointer carry no lock state and an empty [context.Context].
 type checked struct {
-	context.Context             // nil after [checked.Unlock] is called
-	mu              *sync.Mutex // nil if the context does not track a mutex lock state
+	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
-}
-
-func noneChecked() *checked {
-	return &checked{noneCtx, nil, nil}
+	unlocked        bool        // whether [checked.Unlock] was called
 }
 
 func wrapChecked(parent context.Context) *checked {
-	return &checked{parent, nil, nil}
+	return &checked{parent, nil, nil, false}
 }
 
 func lockChecked(parent *checked, mu *sync.Mutex) *checked {
-	checkLockArgs(parent, mu)
+	panicIfNil(mu)
 	if parentLockCtx, 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")
 		}
-		return &checked{parent, mu, parentLockCtx}
+		return &checked{parent, mu, parentLockCtx, false}
 	}
 	mu.Lock()
-	return &checked{parent, mu, nil}
+	return &checked{parent, mu, nil, false}
+}
+
+func (c *checked) Deadline() (deadline time.Time, ok bool) {
+	c.panicIfUnlocked()
+	if c == nil || c.Context == nil {
+		return time.Time{}, false
+	}
+	return c.Context.Deadline()
+}
+
+func (c *checked) Done() <-chan struct{} {
+	c.panicIfUnlocked()
+	if c == nil || c.Context == nil {
+		return nil
+	}
+	return c.Context.Done()
+}
+
+func (c *checked) Err() error {
+	c.panicIfUnlocked()
+	if c == nil || c.Context == nil {
+		return nil
+	}
+	return c.Context.Err()
 }
 
 func (c *checked) Value(key any) any {
-	if c.Context == nil {
-		panic("use of context after unlock")
+	c.panicIfUnlocked()
+	if c == nil {
+		return nil
 	}
 	if key == ctxKeyOf(c.mu) {
 		return c
@@ -69,16 +90,19 @@ func (c *checked) Value(key any) any {
 
 func (c *checked) Unlock() {
 	switch {
-	case c.Context == nil:
+	case c == nil:
+		// No-op; zero context.
+		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].
 	case c.parent == nil:
 		// We own the lock; let's unlock it.
-		// This panics if the mutex is already unlocked.
+		// This triggers a fatal error if the mutex is already unlocked.
 		c.mu.Unlock()
-	case c.parent.Context == nil:
+	case c.parent.unlocked:
 		// The parent context is already unlocked.
 		// The mutex may or may not be locked;
 		// something else may have already locked it.
@@ -92,25 +116,26 @@ func (c *checked) Unlock() {
 	default:
 		// No-op; a parent or ancestor will handle unlocking.
 	}
-	c.Context = nil
+	c.unlocked = true // mark this context as unlocked
 }
 
-func checkLockArgs[T interface {
-	context.Context
-	comparable
-}](parent T, mu *sync.Mutex) {
-	var zero T
-	if parent == zero {
-		panic("nil parent context")
+func (c *checked) panicIfUnlocked() {
+	if c != nil && c.unlocked {
+		panic("use of context after unlock")
 	}
-	if mu == nil {
-		panic(fmt.Sprintf("nil %T", mu))
+}
+
+func panicIfNil[T comparable](v T) {
+	if reflect.ValueOf(v).IsNil() {
+		panic(fmt.Sprintf("nil %T", v))
 	}
 }
 
 // unchecked is an implementation of [Context] that trades runtime checks for performance.
+//
+// Its zero value carries no mutex lock state and an empty [context.Context].
 type unchecked struct {
-	context.Context             // always non-nil
+	context.Context             // nil if the context does not carry a [context.Context]
 	mu              *sync.Mutex // non-nil if locked by this context
 }
 
@@ -119,7 +144,6 @@ func wrapUnchecked(parent context.Context) unchecked {
 }
 
 func lockUnchecked(parent unchecked, mu *sync.Mutex) unchecked {
-	checkLockArgs(parent, mu) // this is cheap, so we do it even in the unchecked case
 	if parent.Value(ctxKeyOf(mu)) == nil {
 		mu.Lock()
 	} else {
@@ -128,10 +152,34 @@ func lockUnchecked(parent unchecked, mu *sync.Mutex) unchecked {
 	return unchecked{parent.Context, mu}
 }
 
+func (c unchecked) Deadline() (deadline time.Time, ok bool) {
+	if c.Context == nil {
+		return time.Time{}, false
+	}
+	return c.Context.Deadline()
+}
+
+func (c unchecked) Done() <-chan struct{} {
+	if c.Context == nil {
+		return nil
+	}
+	return c.Context.Done()
+}
+
+func (c unchecked) Err() error {
+	if c.Context == nil {
+		return nil
+	}
+	return c.Context.Err()
+}
+
 func (c unchecked) Value(key any) any {
 	if key == ctxKeyOf(c.mu) {
 		return key
 	}
+	if c.Context == nil {
+		return nil
+	}
 	return c.Context.Value(key)
 }
 

util/ctxlock/ctx_checked.go

@@ -26,7 +26,7 @@ type Context struct {
 // 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{noneChecked()}
+	return Context{}
 }
 
 // Wrap returns a derived [Context] that wraps the provided [context.Context].

util/ctxlock/ctx_test.go

@@ -29,12 +29,12 @@ var (
 		Lock: Lock,
 	}
 	checkedImpl = impl[*checked]{
-		None: noneChecked,
+		None: func() *checked { return nil },
 		Wrap: wrapChecked,
 		Lock: lockChecked,
 	}
 	uncheckedImpl = impl[unchecked]{
-		None: func() unchecked { return noneUnchecked },
+		None: func() unchecked { return unchecked{} },
 		Wrap: wrapUnchecked,
 		Lock: lockUnchecked,
 	}
@@ -158,28 +158,9 @@ func testWrappedLockContext[T ctx](t *testing.T, impl impl[T]) {
 	wantUnlocked(t, &mu) // mu is now unlocked
 }
 
-func TestNilContextAndMutex(t *testing.T) {
-	t.Run("Exported", func(t *testing.T) {
-		testNilContextAndMutex(t, exportedImpl)
-	})
-
-	t.Run("Checked", func(t *testing.T) {
-		testNilContextAndMutex(t, checkedImpl)
-	})
-
-	t.Run("Unchecked", func(t *testing.T) {
-		testNilContextAndMutex(t, uncheckedImpl)
-	})
-}
-
-func testNilContextAndMutex[T ctx](t *testing.T, impl impl[T]) {
-	t.Run("NilContext", func(t *testing.T) {
-		var zero T
-		wantPanic(t, "nil parent context", func() { impl.Lock(zero, &sync.Mutex{}) })
-	})
-	t.Run("NilMutex", func(t *testing.T) {
-		wantPanic(t, "nil *sync.Mutex", func() { impl.Lock(impl.None(), nil) })
-	})
+func TestNilMutex(t *testing.T) {
+	impl := checkedImpl
+	wantPanic(t, "nil *sync.Mutex", func() { impl.Lock(impl.None(), nil) })
 }
 
 func TestUseUnlockedParent_Checked(t *testing.T) {
@@ -205,7 +186,7 @@ func TestUnlockParentFirst_Checked(t *testing.T) {
 	impl := checkedImpl
 
 	var mu sync.Mutex
-	parent := impl.Lock(impl.None(), &mu)
+	parent := impl.Lock(impl.Wrap(context.Background()), &mu)
 	child := impl.Lock(parent, &mu)
 
 	parent.Unlock()      // unlocks mu
@@ -221,9 +202,6 @@ func TestUnlockTwice_Checked(t *testing.T) {
 		wantPanic(t, "already unlocked", ctx.Unlock)
 	}
 
-	t.Run("None", func(t *testing.T) {
-		unlockTwice(t, impl.None())
-	})
 	t.Run("Wrapped", func(t *testing.T) {
 		unlockTwice(t, impl.Wrap(context.Background()))
 	})

util/ctxlock/ctx_unchecked.go

@@ -18,7 +18,7 @@ type Context struct {
 }
 
 func None() Context {
-	return Context{noneUnchecked}
+	return Context{}
 }
 
 func Wrap(parent context.Context) Context {

util/ctxlock/doc_test.go

@@ -74,3 +74,16 @@ func ExampleContext_twoResources() {
 	fmt.Println(r1.GetFoo(ctxlock.None()))
 	// 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
+}