util/ctxlock: make ctxlock.Lock generic

So that it works with both ctxlock.Context and context.Context without allocating
and without requiring ctxlock.Wrap at callsites.

Updates #12614

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

util/ctxlock/ctx_checked.go

@@ -43,6 +43,9 @@ func Wrap(parent context.Context) Context {
 // 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(parent Context, mu *sync.Mutex) Context {
-	return Context{lockChecked(parent.checked, mu)}
+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_test.go

@@ -26,7 +26,7 @@ var (
 	exportedImpl = impl[Context]{
 		None: None,
 		Wrap: Wrap,
-		Lock: Lock,
+		Lock: Lock[Context],
 	}
 	checkedImpl = impl[*checked]{
 		None: func() *checked { return nil },
@@ -77,6 +77,34 @@ func benchmarkReentrance[T ctx](b *testing.B, impl impl[T]) {
 	}
 }
 
+func BenchmarkGenericLock(b *testing.B) {
+	// Does not allocate with --tags=ts_omit_ctxlock_checks.
+	b.Run("ZeroContext", func(b *testing.B) {
+		var mu sync.Mutex
+		var ctx Context
+		for b.Loop() {
+			parent := Lock(ctx, &mu)
+			func(ctx Context) {
+				child := Lock(ctx, &mu)
+				child.Unlock()
+			}(parent)
+			parent.Unlock()
+		}
+	})
+	b.Run("StdContext", func(b *testing.B) {
+		var mu sync.Mutex
+		ctx := context.Background()
+		for b.Loop() {
+			parent := Lock(ctx, &mu)
+			func(ctx Context) {
+				child := Lock(ctx, &mu)
+				child.Unlock()
+			}(parent)
+			parent.Unlock()
+		}
+	})
+}
+
 func TestHappyPath(t *testing.T) {
 	t.Run("Exported", func(t *testing.T) {
 		testHappyPath(t, exportedImpl)

util/ctxlock/ctx_unchecked.go

@@ -25,6 +25,9 @@ func Wrap(parent context.Context) Context {
 	return Context{wrapUnchecked(parent)}
 }
 
-func Lock(parent Context, mu *sync.Mutex) Context {
-	return Context{lockUnchecked(parent.unchecked, mu)}
+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

@@ -4,10 +4,10 @@
 package ctxlock_test
 
 import (
+	"context"
 	"fmt"
 	"sync"
 
-	"tailscale.com/syncs"
 	"tailscale.com/util/ctxlock"
 )
 
@@ -27,7 +27,8 @@ func (r *Resource) SetFoo(ctx ctxlock.Context, foo string) {
 }
 
 func (r *Resource) GetBar(ctx ctxlock.Context) string {
-	defer ctxlock.Lock(ctx, &r.mu).Unlock()
+	ctx = ctxlock.Lock(ctx, &r.mu)
+	defer ctx.Unlock() // If you prefer it this way.
 	return r.bar
 }
 
@@ -43,6 +44,17 @@ func (r *Resource) WithLock(ctx ctxlock.Context, f func(ctx ctxlock.Context)) {
 	f(ctx) // Call the callback with the new context.
 }
 
+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].
+	// [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()
+	r.foo = foo
+	r.bar = bar
+	return f(muCtx)
+}
+
 func ExampleContext() {
 	var r Resource
 	r.SetFoo(ctxlock.None(), "foo")
@@ -82,3 +94,14 @@ func ExampleContext_zeroValue() {
 	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 {
+		// The r's lock is held, and ctx carries the lock state.
+		return r.GetFoo(ctx) + r.GetBar(ctx)
+	})
+	fmt.Println(result)
+	// Output: foobar
+}