proposal: context: add generic Key type

[dsnet]

The Context API suffers from the lack of type safety since keys and values are both interface{}.

With generics coming in Go 1.18, I propose the addition of:

package context

// Key is a generic key type associated with a specific value type.
type Key[Value any] struct{ name *string }

// NewKey constructs a new context key with an associated value.
// Every key is unique, even if provided the same name.
func NewKey[Value any](name string) Key[Value] {
	return Key[Value]{&name}
}

// WithValue returns a copy of parent in which the value associated with key is val.
//
// It is a type-safe equivalent of WithValue.
func (key Key[Value]) WithValue(parent context.Context, val Value) context.Context {
	return context.WithValue(ctx, key, val)
}

// Value returns the value associated with this context for key,
// additionally reporting whether the value was found.
//
// It is a type-safe equivalent of Context.Value.
func (key Key[Value]) Value(ctx context.Context) (Value, bool) {
	v, ok := ctx.Value(key).(Value)
	return v, ok
}

// String returns the name of the key.
func (key Key[Value]) String() string {
	return *key.name
}

For type safety, the value type is associated with the key type, and type safe variants of Context.Value and WithValue are hung off of Key as methods.

Example usage:

k1 := NewKey[error]("error")
k2 := NewKey[error]("error")
fmt.Println(k1 == k2) // expect false since each key is unique

ctx := context.Background()
ctx = k1.WithValue(ctx, io.EOF)

fmt.Println(k1.Value(ctx)) // EOF true
fmt.Println(k2.Value(ctx)) // <nil> false

\cc @josharian

[ianlancetaylor]

This seems to have the potential to cause confusion if multiple packages use the type context.Key[string]. As with the WithValue method, it seems essential to only instantiate context.Key with a non-builtin type.

[DeedleFake]

@ianlancetaylor

I thought of that, too, but the type-safety that this affords could still be gained by adding two types to it:

type Key[K comparable, V any] struct {
  key K
}

// This could be weird because of type inference not handling `K` due to `V` having to be manual.
func NewKey[K comparable, V any](key K) Key[K, V] {
  return Key[K, V]{key: key}
}

func (k Key[K, V]) WithValue(ctx context.Context, val V) context.Context {
  return WithValue(ctx, k.key, val)
}

func (k Key[K, V]) Value(ctx context.Context) (V, bool) {
  v, ok := ctx.Value(k.key).(V)
  return v, ok
}

Since the standard practice is to just expose access via functions that do this manually for each type anyways, though, I'm not sure that this really gains that much.

[balasanjay]

@dsnet A minor nit:

With an implementation like return context.WithValue(ctx, key, val), its possible to bypass the type-safe API by writing that directly (and therefore to inadvertently introduce non-type-safe writes).

Should we instead implement it like return context.WithValue(ctx, key.name, val)?
That provides a strong guarantee that for any use-cases using the context.Key type, all accesses go through the type-safe functions.

[dsnet]

This seems to have the potential to cause confusion if multiple packages use the type context.Key[string].

Possibly. My hope is that users gets accustomed to the idea that each key returned by NewKey is unique in a similar way to how every error returned by errors.New is unique (i.e., io.EOF != errors.New("EOF")).

Should we instead implement it like return context.WithValue(ctx, key.name, val)?
That provides a strong guarantee that for any use-cases using the context.Key type, all accesses go through the type-safe functions.

Great idea.

[neild]

What about dropping keys entirely and using the type as the key? Something like:

func WithValueOf[T any](parent Context, val T) Context
func ValueOf[T any](ctx Context) (val T, ok bool)

type traceID int
ctx := context.WithValueOf(context.Background(), traceID(42))
id, ok := context.ValueOf[traceID](ctx)

[dsnet]

What about dropping keys entirely and using the type as the key?

I have mixed feelings. I like the simplicity, but it increases the chance of two packages both operating on a common type (i.e., a string value) where they work in isolation, but when used together they conflict and fight with one another.

[earthboundkid]

It's not clear to me why NewKey takes a string. Why not something like?

-- mypackage/context.go --
package mypackage

var (
    LoggerCtx = context.NewValue[Logger]()
    TraceIDCtx = context.NewValue[TraceID]()
)

func DoSomething(ctx context.Context) {
   logger, ok := LoggerCtx.Value(ctx)
   if ok {
       traceid, _ := TraceIDCtx.Value(ctx)
       logger("wow, look at this trace ID", traceid)
   }
}

[dsnet]

It's not clear to me why NewKey takes a string.

The provided string name is purely for debugging purposes. I don't feel strongly about having it.

[earthboundkid]

Okay, I see. I guess &name guaranteed to be unique because it's a closed over value, but it's not essential to the design to use it as the key. Too clever for me. :-)

[deanveloper]

Keeping NewKey also adds the benefit of an additional constraint to be passed in to the value-retrieval function. this helps with type-safety, especially when there may be more than one value for a given type:

type User struct { ... }
var SendingUserCtx = context.NewKey[User]("sending user")
var TargetUserCtx = context.NewKey[User]("target user")

func abc(ctx context.Context) {
    sender := SendingUserCtx.Value(ctx)
}

// vs

type User struct { ... }
type SendingUserCtx User
type TargetUserCtx User

func abc(ctx context.Context) {
    // possible confusion: may accidentally provide the wrong key here
    sender := context.ValueOf[User](ctx)
}

[deanveloper]

Something nice to mention is that this addresses 2 of the 4 things mentioned in #28342 (namely type-safety and documentation)

[earthboundkid]

The more I look at @dsnet's proposed implementation, the more I like it. My only suggestion is to change WithValue to panic if name is nil, to prevent accidental misuse of the zero value.

[rsc]

On hold for understanding generics better.

[rsc]

Placed on hold.
— rsc for the proposal review group

[xaionaro]

What about dropping keys entirely and using the type as the key?

I have mixed feelings. I like the simplicity, but it increases the chance of two packages both operating on a common type (i.e., a string value) where they work in isolation, but when used together they conflict and fight with one another.

Well, in the assumption of possibility to break some backward compatibility in Go2 it does not seem bad to just forbid non-named types in WithValue and Value (a compile-time error if possible or a panic + go lint if not possible).

[cpl]

There should be one more aspect considered. Right now a common pattern is to have the context key private within the package, and then provide a public function which may do more than just check the given context and return (Value, bool).

For example this function which will return a new (or a default) logger if one isn't set on the context. The idea is, get value from context functions sometimes have to be a bit smarter.

type ctxKeyLogger struct{}

func CtxLogger(ctx context.Context) Logger {
	logger, ok := ctx.Value(ctxKeyLogger{})
	if !ok {
		logger = NewLogger()
	}
	
	return logger
}

What I was thinking of is having something along the lines of:

type Key[T comparable] struct {
	valueFunc keyValueFunc[T]
}

type keyValueFunc[T comparable] func(ctx context.Context, k Key[T]) (T, bool)

func defaultValueFunc[T comparable](ctx context.Context, k Key[T]) (T, bool) {
	value, ok := ctx.Value(k).(T)
	return value, ok
}

func NewKey[T comparable](f keyValueFunc[T]) Key[T] {
	if f == nil {
		f = defaultValueFunc[T]
	}

	return Key[T]{valueFunc: f}
}

func (k Key[T]) Value(ctx context.Context) (T, bool) {
	return k.valueFunc(ctx, k)
}

func (k Key[T]) WithValue(ctx context.Context, value T) context.Context {
	return context.WithValue(ctx, k, value)
}

I've played around with an interface version of Key too:

type Key[T comparable] interface {
	Value(ctx context.Context) (T, bool)
	WithValue(ctx context.Context, value T) context.Context
}

Honestly, from both my ideas, I am not 100% sold on them myself, but maybe it gives someone a spark of imagination!

[jackielii]

@dsnet Could we create a package for this? I think @ianlancetaylor has a point in #57025 (comment):

I don't see a lot of support for this proposal, and it can be implemented outside of the standard library. Perhaps it should be an external package for now, and we can consider adopting it into the standard library if it sees wide use.

I've been copy and pasting your sample code for many times now. Creating a package so that we can re-use and see how many people does the same seems useful to me.

[dsnet]

There's https://pkg.go.dev/tailscale.com/util/ctxkey, but we provide no guarantees about API compatibility.