ctxkey: add type-safe context keys (#123)

Author: astrophena

.devtools/config.txtar

@@ -1,5 +1,7 @@
 -- copyright/exclusions.json --
 [
+  "ctxkey/ctxkey.go",
+  "ctxkey/ctxkey_test.go",
   "web/internal/hashfs/*.go",
   "web/internal/unionfs/*.go",
   "rr/*.go",

cli/cli.go

@@ -82,6 +82,7 @@ import (
 	"strings"
 	"time"
 
+	"go.astrophena.name/base/ctxkey"
 	"go.astrophena.name/base/logger"
 	"go.astrophena.name/base/syncx"
 	"go.astrophena.name/base/version"
@@ -179,24 +180,19 @@ func (f AppFunc) Run(ctx context.Context) error {
 	return f(ctx)
 }
 
-type ctxKey int
-
-var envKey ctxKey
+var envKey = ctxkey.New[*Env]("cli.Env", nil)
 
 // GetEnv retrieves the application's environment from a context.
 // If the context has no environment, it returns one based on the current OS.
 func GetEnv(ctx context.Context) *Env {
-	e, ok := ctx.Value(envKey).(*Env)
-	if !ok {
-		return OSEnv()
+	if e, ok := envKey.ValueOk(ctx); ok {
+		return e
 	}
-	return e
+	return OSEnv()
 }
 
 // WithEnv returns a new context that carries the provided [Env].
-func WithEnv(ctx context.Context, e *Env) context.Context {
-	return context.WithValue(ctx, envKey, e)
-}
+func WithEnv(ctx context.Context, e *Env) context.Context { return envKey.WithValue(ctx, e) }
 
 // Env encapsulates the application's environment, including arguments,
 // standard I/O streams, and environment variables.

ctxkey/LICENSE

@@ -0,0 +1,28 @@
+BSD 3-Clause License
+
+Copyright (c) 2020 Tailscale Inc & contributors.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

ctxkey/ctxkey.go

@@ -0,0 +1,135 @@
+// Copyright (c) Tailscale Inc & contributors
+// SPDX-License-Identifier: BSD-3-Clause
+
+// ctxkey provides type-safe key-value pairs for use with [context.Context].
+//
+// Example usage:
+//
+//	// Create a context key.
+//	var TimeoutKey = ctxkey.New("mapreduce.Timeout", 5*time.Second)
+//
+//	// Store a context value.
+//	ctx = mapreduce.TimeoutKey.WithValue(ctx, 10*time.Second)
+//
+//	// Load a context value.
+//	timeout := mapreduce.TimeoutKey.Value(ctx)
+//	... // use timeout of type time.Duration
+//
+// This is inspired by https://go.dev/issue/49189.
+package ctxkey
+
+import (
+	"context"
+	"fmt"
+	"reflect"
+)
+
+// Key is a generic key type associated with a specific value type.
+//
+// A zero Key is valid where the Value type itself is used as the context key.
+// This pattern should only be used with locally declared Go types,
+// otherwise different packages risk producing key conflicts.
+//
+// Example usage:
+//
+//	type peerInfo struct { ... }           // peerInfo is a locally declared type
+//	var peerInfoKey ctxkey.Key[peerInfo]
+//	ctx = peerInfoKey.WithValue(ctx, info) // store a context value
+//	info = peerInfoKey.Value(ctx)          // load a context value
+type Key[Value any] struct {
+	name   *stringer[string]
+	defVal *Value
+}
+
+// New constructs a new context key with an associated value type
+// where the default value for an unpopulated value is the provided value.
+//
+// The provided name is an arbitrary name only used for human debugging.
+// As a convention, it is recommended that the name be the dot-delimited
+// combination of the package name of the caller with the variable name.
+// If the name is not provided, then the name of the Value type is used.
+// Every key is unique, even if provided the same name.
+//
+// Example usage:
+//
+//	package mapreduce
+//	var NumWorkersKey = ctxkey.New("mapreduce.NumWorkers", runtime.NumCPU())
+func New[Value any](name string, defaultValue Value) Key[Value] {
+	// Allocate a new stringer to ensure that every invocation of New
+	// creates a universally unique context key even for the same name
+	// since newly allocated pointers are globally unique within a process.
+	key := Key[Value]{name: new(stringer[string])}
+	if name == "" {
+		name = reflect.TypeFor[Value]().String()
+	}
+	key.name.v = name
+	if v := reflect.ValueOf(defaultValue); v.IsValid() && !v.IsZero() {
+		key.defVal = &defaultValue
+	}
+	return key
+}
+
+// contextKey returns the context key to use.
+func (key Key[Value]) contextKey() any {
+	if key.name == nil {
+		// Use the reflect.Type of the Value (implies key not created by New).
+		return reflect.TypeFor[Value]()
+	} else {
+		// Use the name pointer directly (implies key created by New).
+		return key.name
+	}
+}
+
+// WithValue returns a copy of parent in which the value associated with key is val.
+//
+// It is a type-safe equivalent of [context.WithValue].
+func (key Key[Value]) WithValue(parent context.Context, val Value) context.Context {
+	return context.WithValue(parent, key.contextKey(), stringer[Value]{val})
+}
+
+// ValueOk returns the value in the context associated with this key
+// and also reports whether it was present.
+// If the value is not present, it returns the default value.
+func (key Key[Value]) ValueOk(ctx context.Context) (v Value, ok bool) {
+	vv, ok := ctx.Value(key.contextKey()).(stringer[Value])
+	if !ok && key.defVal != nil {
+		vv.v = *key.defVal
+	}
+	return vv.v, ok
+}
+
+// Value returns the value in the context associated with this key.
+// If the value is not present, it returns the default value.
+func (key Key[Value]) Value(ctx context.Context) (v Value) {
+	v, _ = key.ValueOk(ctx)
+	return v
+}
+
+// Has reports whether the context has a value for this key.
+func (key Key[Value]) Has(ctx context.Context) (ok bool) {
+	_, ok = key.ValueOk(ctx)
+	return ok
+}
+
+// String returns the name of the key.
+func (key Key[Value]) String() string {
+	if key.name == nil {
+		return reflect.TypeFor[Value]().String()
+	}
+	return key.name.String()
+}
+
+// stringer implements [fmt.Stringer] on a generic T.
+//
+// This assists in debugging such that printing a context prints key and value.
+// Note that the [context] package lacks a dependency on [reflect],
+// so it cannot print arbitrary values. By implementing [fmt.Stringer],
+// we functionally teach a context how to print itself.
+//
+// Wrapping values within a struct has an added bonus that interface kinds
+// are properly handled. Without wrapping, we would be unable to distinguish
+// between a nil value that was explicitly set or not.
+// However, the presence of a stringer indicates an explicit nil value.
+type stringer[T any] struct{ v T }
+
+func (v stringer[T]) String() string { return fmt.Sprint(v.v) }

ctxkey/ctxkey_test.go

@@ -0,0 +1,107 @@
+// Copyright (c) Tailscale Inc & contributors
+// SPDX-License-Identifier: BSD-3-Clause
+
+package ctxkey
+
+import (
+	"fmt"
+	"io"
+	"regexp"
+	"testing"
+	"time"
+
+	"go.astrophena.name/base/testutil"
+)
+
+func TestKey(t *testing.T) {
+	ctx := t.Context()
+
+	// Test keys with the same name as being distinct.
+	k1 := New("same.Name", "")
+	testutil.AssertEqual(t, k1.String(), "same.Name")
+	k2 := New("same.Name", "")
+	testutil.AssertEqual(t, k2.String(), "same.Name")
+	testutil.AssertEqual(t, k1 == k2, false)
+	ctx = k1.WithValue(ctx, "hello")
+	testutil.AssertEqual(t, k1.Has(ctx), true)
+	testutil.AssertEqual(t, k1.Value(ctx), "hello")
+	testutil.AssertEqual(t, k2.Has(ctx), false)
+	testutil.AssertEqual(t, k2.Value(ctx), "")
+	ctx = k2.WithValue(ctx, "goodbye")
+	testutil.AssertEqual(t, k1.Has(ctx), true)
+	testutil.AssertEqual(t, k1.Value(ctx), "hello")
+	testutil.AssertEqual(t, k2.Has(ctx), true)
+	testutil.AssertEqual(t, k2.Value(ctx), "goodbye")
+
+	// Test default value.
+	k3 := New("mapreduce.Timeout", time.Hour)
+	testutil.AssertEqual(t, k3.Has(ctx), false)
+	testutil.AssertEqual(t, k3.Value(ctx), time.Hour)
+	ctx = k3.WithValue(ctx, time.Minute)
+	testutil.AssertEqual(t, k3.Has(ctx), true)
+	testutil.AssertEqual(t, k3.Value(ctx), time.Minute)
+
+	// Test incomparable value.
+	k4 := New("slice", []int(nil))
+	testutil.AssertEqual(t, k4.Has(ctx), false)
+	testutil.AssertEqual(t, k4.Value(ctx), []int(nil))
+	ctx = k4.WithValue(ctx, []int{1, 2, 3})
+	testutil.AssertEqual(t, k4.Has(ctx), true)
+	testutil.AssertEqual(t, k4.Value(ctx), []int{1, 2, 3})
+
+	// Accessors should be allocation free.
+	testutil.AssertEqual(t, testing.AllocsPerRun(100, func() {
+		k1.Value(ctx)
+		k1.Has(ctx)
+		k1.ValueOk(ctx)
+	}), 0.0)
+
+	// Test keys that are created without New.
+	var k5 Key[string]
+	testutil.AssertEqual(t, k5.String(), "string")
+	testutil.AssertEqual(t, k1 == k5, false) // should be different from key created by New
+	testutil.AssertEqual(t, k5.Has(ctx), false)
+	ctx = k5.WithValue(ctx, "fizz")
+	testutil.AssertEqual(t, k5.Value(ctx), "fizz")
+	var k6 Key[string]
+	testutil.AssertEqual(t, k6.String(), "string")
+	testutil.AssertEqual(t, k5 == k6, true)
+	testutil.AssertEqual(t, k6.Has(ctx), true)
+	ctx = k6.WithValue(ctx, "fizz")
+
+	// Test interface value types.
+	var k7 Key[any]
+	testutil.AssertEqual(t, k7.Has(ctx), false)
+	ctx = k7.WithValue(ctx, "whatever")
+	testutil.AssertEqual(t, k7.Value(ctx), "whatever")
+	ctx = k7.WithValue(ctx, []int{1, 2, 3})
+	testutil.AssertEqual(t, k7.Value(ctx), []int{1, 2, 3})
+	ctx = k7.WithValue(ctx, nil)
+	testutil.AssertEqual(t, k7.Has(ctx), true)
+	testutil.AssertEqual(t, k7.Value(ctx), nil)
+	k8 := New[error]("error", io.EOF)
+	testutil.AssertEqual(t, k8.Has(ctx), false)
+	testutil.AssertEqual(t, k8.Value(ctx), io.EOF)
+	ctx = k8.WithValue(ctx, nil)
+	testutil.AssertEqual(t, k8.Value(ctx), nil)
+	testutil.AssertEqual(t, k8.Has(ctx), true)
+	err := fmt.Errorf("read error: %w", io.ErrUnexpectedEOF)
+	ctx = k8.WithValue(ctx, err)
+	testutil.AssertEqual(t, k8.Value(ctx), err)
+	testutil.AssertEqual(t, k8.Has(ctx), true)
+}
+
+func TestStringer(t *testing.T) {
+	ctx := t.Context()
+	assertMatches(t, fmt.Sprint(New("foo.Bar", "").WithValue(ctx, "baz")), regexp.MustCompile("foo.Bar.*baz"))
+	assertMatches(t, fmt.Sprint(New("", []int{}).WithValue(ctx, []int{1, 2, 3})), regexp.MustCompile(fmt.Sprintf("%[1]T.*%[1]v", []int{1, 2, 3})))
+	assertMatches(t, fmt.Sprint(New("", 0).WithValue(ctx, 5)), regexp.MustCompile("int.*5"))
+	assertMatches(t, fmt.Sprint(Key[time.Duration]{}.WithValue(ctx, time.Hour)), regexp.MustCompile(fmt.Sprintf("%[1]T.*%[1]v", time.Hour)))
+}
+
+func assertMatches(t *testing.T, got string, re *regexp.Regexp) {
+	t.Helper()
+	if !re.MatchString(got) {
+		t.Fatalf("value does not match regexp:\ngot:    %q\nregexp: %q", got, re)
+	}
+}

logger/logger.go

@@ -10,11 +10,11 @@ import (
 	"io"
 	"log/slog"
 	"sync"
-)
 
-type ctxKey string
+	"go.astrophena.name/base/ctxkey"
+)
 
-const loggerKey ctxKey = "logger"
+var loggerKey = ctxkey.New("logger.Logger", defaultLogger)
 
 // multiHandler fans out log records to multiple handlers.
 type multiHandler struct {
@@ -135,20 +135,13 @@ func newDefaultLogger() *Logger {
 }
 
 // Put returns a new context with the provided [Logger].
-func Put(ctx context.Context, l *Logger) context.Context {
-	return context.WithValue(ctx, loggerKey, l)
-}
+func Put(ctx context.Context, l *Logger) context.Context { return loggerKey.WithValue(ctx, l) }
 
 // Get retrieves the [Logger] from the context.
 //
 // If the context has no [Logger], it returns a default [Logger] that discards all
 // messages.
-func Get(ctx context.Context) *Logger {
-	if l, ok := ctx.Value(loggerKey).(*Logger); ok {
-		return l
-	}
-	return defaultLogger
-}
+func Get(ctx context.Context) *Logger { return loggerKey.Value(ctx) }
 
 // IsDefault returns true if l is the default [Logger].
 func IsDefault(l *Logger) bool { return l == defaultLogger }
@@ -157,12 +150,7 @@ func IsDefault(l *Logger) bool { return l == defaultLogger }
 //
 // If the context has no [Logger], it returns a [slog.LevelVar] for a default
 // [Logger].
-func LevelVar(ctx context.Context) *slog.LevelVar {
-	if l, ok := ctx.Value(loggerKey).(*Logger); ok {
-		return l.Level
-	}
-	return defaultLogger.Level
-}
+func LevelVar(ctx context.Context) *slog.LevelVar { return loggerKey.Value(ctx).Level }
 
 // Debug logs a debug message.
 func Debug(ctx context.Context, msg string, attrs ...slog.Attr) {

web/debug.go

@@ -180,7 +180,7 @@ func (s *Server) xffDebugHandler() http.Handler {
 			parts = append(parts, item)
 		}
 
-		connNetwork, _ := r.Context().Value(connNetworkContextKey).(string)
+		connNetwork := connNetworkContextKey.Value(r.Context())
 		resp := xffDebugResponse{
 			RemoteAddr:                 r.RemoteAddr,
 			RemoteHost:                 host,