-
Notifications
You must be signed in to change notification settings - Fork 1.3k
/
key.go
135 lines (123 loc) · 4.54 KB
/
key.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
// Copyright (c) Tailscale Inc & AUTHORS
// 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) }