mctx: refactor to no longer have parent/child logic

This commit is contained in:
Brian Picciano 2019-06-15 17:28:29 -06:00
parent c98f154992
commit 467bcbe52d
5 changed files with 49 additions and 598 deletions

View File

@ -4,15 +4,12 @@ import (
"context" "context"
"fmt" "fmt"
"sort" "sort"
"strings"
) )
// Annotation describes the annotation of a key/value pair made on a Context via // Annotation describes the annotation of a key/value pair made on a Context via
// the Annotate call. The Path field is the Path of the Context on which the // the Annotate call.
// call was made.
type Annotation struct { type Annotation struct {
Key, Value interface{} Key, Value interface{}
Path []string
} }
type annotation struct { type annotation struct {
@ -23,11 +20,7 @@ type annotation struct {
type annotationKey int type annotationKey int
// Annotate takes in one or more key/value pairs (kvs' length must be even) and // Annotate takes in one or more key/value pairs (kvs' length must be even) and
// returns a Context carrying them. Annotations only exist on the local level, // returns a Context carrying them.
// i.e. a child and parent share different annotation namespaces.
//
// NOTE that annotations are preserved across NewChild calls, but are keyed
// based on the passed in key _and_ the Context's Path.
func Annotate(ctx context.Context, kvs ...interface{}) context.Context { func Annotate(ctx context.Context, kvs ...interface{}) context.Context {
if len(kvs)%2 > 0 { if len(kvs)%2 > 0 {
panic("kvs being passed to mctx.Annotate must have an even number of elements") panic("kvs being passed to mctx.Annotate must have an even number of elements")
@ -43,13 +36,9 @@ func Annotate(ctx context.Context, kvs ...interface{}) context.Context {
if prev != nil { if prev != nil {
root = prev.root root = prev.root
} }
path := Path(ctx)
for i := 0; i < len(kvs); i += 2 { for i := 0; i < len(kvs); i += 2 {
curr = &annotation{ curr = &annotation{
Annotation: Annotation{ Annotation: Annotation{Key: kvs[i], Value: kvs[i+1]},
Key: kvs[i], Value: kvs[i+1],
Path: path,
},
prev: prev, prev: prev,
} }
if root == nil { if root == nil {
@ -81,11 +70,7 @@ func Annotations(ctx context.Context) AnnotationSet {
if a == nil { if a == nil {
return nil return nil
} }
type mKey struct { m := map[interface{}]bool{}
pathHash string
key interface{}
}
m := map[mKey]bool{}
var aa AnnotationSet var aa AnnotationSet
for { for {
@ -93,33 +78,29 @@ func Annotations(ctx context.Context) AnnotationSet {
break break
} }
k := mKey{pathHash: pathHash(a.Path), key: a.Key} if m[a.Key] {
if m[k] {
a = a.prev a = a.prev
continue continue
} }
aa = append(aa, a.Annotation) aa = append(aa, a.Annotation)
m[k] = true m[a.Key] = true
a = a.prev a = a.prev
} }
return aa return aa
} }
// StringMapByPath is similar to StringMap, but it first maps each annotation // StringMap formats each of the Annotations into strings using fmt.Sprint. If
// datum by its path. // any two keys format to the same string, then type information will be
func (aa AnnotationSet) StringMapByPath() map[string]map[string]string { // prefaced to each one.
func (aa AnnotationSet) StringMap() map[string]string {
type mKey struct { type mKey struct {
str string str string
path string
typ string typ string
} }
m := map[mKey][]Annotation{} m := map[mKey][]Annotation{}
for _, a := range aa { for _, a := range aa {
k := mKey{ k := mKey{str: fmt.Sprint(a.Key)}
str: fmt.Sprint(a.Key),
path: "/" + strings.Join(a.Path, "/"),
}
m[k] = append(m[k], a) m[k] = append(m[k], a)
} }
@ -150,75 +131,12 @@ func (aa AnnotationSet) StringMapByPath() map[string]map[string]string {
} }
} }
outM := map[string]map[string]string{}
for k, annotations := range m {
a := annotations[0]
if outM[k.path] == nil {
outM[k.path] = map[string]string{}
}
kStr := k.str
if k.typ != "" {
kStr += "(" + k.typ + ")"
}
outM[k.path][kStr] = fmt.Sprint(a.Value)
}
return outM
}
// StringMap formats each of the Annotations into strings using fmt.Sprint. If
// any two keys format to the same string, then path information will be
// prefaced to each one. If they still format to the same string, then type
// information will be prefaced to each one.
func (aa AnnotationSet) StringMap() map[string]string {
type mKey struct {
str string
path string
typ string
}
m := map[mKey][]Annotation{}
for _, a := range aa {
k := mKey{str: fmt.Sprint(a.Key)}
m[k] = append(m[k], a)
}
nextK := func(k mKey, a Annotation) mKey {
if k.path == "" {
k.path = "/" + strings.Join(a.Path, "/")
} else if k.typ == "" {
k.typ = fmt.Sprintf("%T", a.Key)
} else {
panic(fmt.Sprintf("mKey %#v is somehow conflicting with another", k))
}
return k
}
for {
var any bool
for k, annotations := range m {
if len(annotations) == 1 {
continue
}
any = true
for _, a := range annotations {
k2 := nextK(k, a)
m[k2] = append(m[k2], a)
}
delete(m, k)
}
if !any {
break
}
}
outM := map[string]string{} outM := map[string]string{}
for k, annotations := range m { for k, annotations := range m {
a := annotations[0] a := annotations[0]
kStr := k.str kStr := k.str
if k.path != "" {
kStr += "(" + k.path + ")"
}
if k.typ != "" { if k.typ != "" {
kStr += "(" + k.typ + ")" kStr = k.typ + "(" + kStr + ")"
} }
outM[kStr] = fmt.Sprint(a.Value) outM[kStr] = fmt.Sprint(a.Value)
} }
@ -277,10 +195,9 @@ func mergeAnnotations(ctxA, ctxB context.Context) context.Context {
// MergeAnnotations sequentially merges the annotation data of the passed in // MergeAnnotations sequentially merges the annotation data of the passed in
// Contexts into the first passed in one. Data from a Context overwrites // Contexts into the first passed in one. Data from a Context overwrites
// overlapping data on all passed in Contexts to the left of it (keeping in mind // overlapping data on all passed in Contexts to the left of it. All other
// that two Annotations must share the same Key _and_ Path to overlap). All // aspects of the first Context remain the same, and that Context is returned
// other aspects of the first Context remain the same, and that Context is // with the new set of Annotation data.
// returned with the new set of Annotation data.
// //
// NOTE this will panic if no Contexts are passed in. // NOTE this will panic if no Contexts are passed in.
func MergeAnnotations(ctxs ...context.Context) context.Context { func MergeAnnotations(ctxs ...context.Context) context.Context {

View File

@ -8,30 +8,16 @@ import (
) )
func TestAnnotate(t *T) { func TestAnnotate(t *T) {
parent := context.Background() ctx := context.Background()
parent = Annotate(parent, "a", "foo") ctx = Annotate(ctx, "a", "foo")
parent = Annotate(parent, "b", "bar") ctx = Annotate(ctx, "b", "bar")
ctx = Annotate(ctx, "b", "BAR")
child := NewChild(parent, "child") annotations := Annotations(ctx)
child = Annotate(child, "a", "Foo")
child = Annotate(child, "a", "FOO")
child = Annotate(child, "c", "BAZ")
parent = WithChild(parent, child)
parentAnnotations := Annotations(parent)
childAnnotations := Annotations(child)
massert.Require(t, massert.Require(t,
massert.Length(parentAnnotations, 2), massert.Length(annotations, 2),
massert.HasValue(parentAnnotations, Annotation{Key: "a", Value: "foo"}), massert.HasValue(annotations, Annotation{Key: "a", Value: "foo"}),
massert.HasValue(parentAnnotations, Annotation{Key: "b", Value: "bar"}), massert.HasValue(annotations, Annotation{Key: "b", Value: "BAR"}),
massert.Length(childAnnotations, 4),
massert.HasValue(childAnnotations, Annotation{Key: "a", Value: "foo"}),
massert.HasValue(childAnnotations, Annotation{Key: "b", Value: "bar"}),
massert.HasValue(childAnnotations,
Annotation{Key: "a", Path: []string{"child"}, Value: "FOO"}),
massert.HasValue(childAnnotations,
Annotation{Key: "c", Path: []string{"child"}, Value: "BAZ"}),
) )
} }
@ -39,32 +25,19 @@ func TestAnnotationsStringMap(t *T) {
type A int type A int
type B int type B int
aa := AnnotationSet{ aa := AnnotationSet{
{Key: 0, Path: nil, Value: "zero"}, {Key: 0, Value: "zero"},
{Key: 1, Path: nil, Value: "one"}, {Key: 1, Value: "one"},
{Key: 1, Path: []string{"foo"}, Value: "ONE"}, {Key: A(2), Value: "two"},
{Key: A(2), Path: []string{"foo"}, Value: "two"}, {Key: B(2), Value: "TWO"},
{Key: B(2), Path: []string{"foo"}, Value: "TWO"},
} }
massert.Require(t, massert.Require(t,
massert.Equal(map[string]string{ massert.Equal(map[string]string{
"0": "zero",
"1(/)": "one",
"1(/foo)": "ONE",
"2(/foo)(mctx.A)": "two",
"2(/foo)(mctx.B)": "TWO",
}, aa.StringMap()),
massert.Equal(map[string]map[string]string{
"/": {
"0": "zero", "0": "zero",
"1": "one", "1": "one",
}, "mctx.A(2)": "two",
"/foo": { "mctx.B(2)": "TWO",
"1": "ONE", }, aa.StringMap()),
"2(mctx.A)": "two",
"2(mctx.B)": "TWO",
},
}, aa.StringMapByPath()),
) )
} }

View File

@ -1,298 +0,0 @@
// Package mctx extends the builtin context package to organize Contexts into a
// hierarchy.
//
// All functions and methods in this package are thread-safe unless otherwise
// noted.
//
// Parents and children
//
// Each node in the hierarchy is given a name and is aware of all of its
// ancestors. The sequence of ancestor's names, ending in the node's name, is
// called its "path". For example:
//
// ctx := context.Background()
// ctxA := mctx.NewChild(ctx, "A")
// ctxB := mctx.NewChild(ctx, "B")
// fmt.Printf("ctx:%#v\n", mctx.Path(ctx)) // prints "ctx:[]string(nil)"
// fmt.Printf("ctxA:%#v\n", mctx.Path(ctxA)) // prints "ctxA:[]string{"A"}
// fmt.Printf("ctxB:%#v\n", mctx.Path(ctxB)) // prints "ctxA:[]string{"A", "B"}
//
// WithChild can be used to incorporate a child into its parent, making the
// parent's children iterable on it:
//
// ctx := context.Background()
// ctxA1 := mctx.NewChild(ctx, "A1")
// ctxA2 := mctx.NewChild(ctx, "A2")
// ctx = mctx.WithChild(ctx, ctxA1)
// ctx = mctx.WithChild(ctx, ctxA2)
// for _, childCtx := range mctx.Children(ctx) {
// fmt.Printf("%q\n", mctx.Name(childCtx)) // prints "A1" then "A2"
// }
//
// Key/Value
//
// The context's key/value namespace is split into two: a space local to a node,
// not inherited from its parent or inheritable by its children
// (WithLocalValue), and the original one which comes with the builtin context
// package (context.WithValue):
//
// ctx := context.Background()
// ctx = context.WithValue(ctx, "inheritableKey", "foo")
// ctx = mctx.WithLocalValue(ctx, "localKey", "bar")
// childCtx := mctx.NewChild(ctx, "child")
//
// // ctx.Value("inheritableKey") == "foo"
// // child.Value("inheritableKey") == "foo"
// // mctx.LocalValue(ctx, "localKey") == "bar"
// // mctx.LocalValue(child, "localKey") == nil
//
// Annotations
//
// Annotations are a special case of local key/values, where the data being
// stored is specifically runtime metadata which would be useful for logging,
// error output, etc... Annotation data might include an IP address of a
// connected client, a userID the client has authenticated as, the primary key
// of a row in a database being queried, etc...
//
// Annotations are always tied to the path of the node they were set on, so that
// even when the annotations of two contexts are merged the annotation data will
// not overlap unless the contexts have the same path.
package mctx
import (
"context"
"crypto/sha256"
"encoding/hex"
"fmt"
)
////////////////////////////////////////////////////////////////////////////////
type ancestryKey int // 0 -> children, 1 -> parent, 2 -> path
const (
ancestryKeyChildren ancestryKey = iota
ancestryKeyChildrenMap
ancestryKeyParent
ancestryKeyPath
)
// Child returns the Context of the given name which was added to parent via
// WithChild, or nil if no Context of that name was ever added.
func Child(parent context.Context, name string) context.Context {
childrenMap, _ := parent.Value(ancestryKeyChildrenMap).(map[string]int)
if len(childrenMap) == 0 {
return nil
}
i, ok := childrenMap[name]
if !ok {
return nil
}
return parent.Value(ancestryKeyChildren).([]context.Context)[i]
}
// Children returns all children of this Context which have been added by
// WithChild, in the order they were added. If this Context wasn't produced by
// WithChild then this returns nil.
func Children(parent context.Context) []context.Context {
children, _ := parent.Value(ancestryKeyChildren).([]context.Context)
return children
}
func childrenCP(parent context.Context) ([]context.Context, map[string]int) {
children := Children(parent)
// plus 1 because this is most commonly used in WithChild, which will append
// to it. At any rate it doesn't hurt anything.
outChildren := make([]context.Context, len(children), len(children)+1)
copy(outChildren, children)
childrenMap, _ := parent.Value(ancestryKeyChildrenMap).(map[string]int)
outChildrenMap := make(map[string]int, len(childrenMap)+1)
for name, i := range childrenMap {
outChildrenMap[name] = i
}
return outChildren, outChildrenMap
}
// parentOf returns the Context from which this one was generated via NewChild.
// Returns nil if this Context was not generated via NewChild.
//
// This is kept private because the behavior is a bit confusing. This will
// return the Context which was passed into NewChild, but users would probably
// expect it to return the one from WithChild if they were to call this.
func parentOf(ctx context.Context) context.Context {
parent, _ := ctx.Value(ancestryKeyParent).(context.Context)
return parent
}
// Path returns the sequence of names which were used to produce this Context
// via the NewChild function. If this Context wasn't produced by NewChild then
// this returns nil.
func Path(ctx context.Context) []string {
path, _ := ctx.Value(ancestryKeyPath).([]string)
return path
}
func pathCP(ctx context.Context) []string {
path := Path(ctx)
// plus 1 because this is most commonly used in NewChild, which will append
// to it. At any rate it doesn't hurt anything.
outPath := make([]string, len(path), len(path)+1)
copy(outPath, path)
return outPath
}
func pathHash(path []string) string {
pathHash := sha256.New()
for _, pathEl := range path {
fmt.Fprintf(pathHash, "%q.", pathEl)
}
return hex.EncodeToString(pathHash.Sum(nil))
}
// Name returns the name this Context was created with via NewChild, or false if
// this Context was not created via NewChild.
func Name(ctx context.Context) (string, bool) {
path := Path(ctx)
if len(path) == 0 {
return "", false
}
return path[len(path)-1], true
}
// NewChild creates and returns a new Context based off of the parent one. The
// child will have a path which is the parent's path appended with the given
// name. In order for the parent to "see" the child (via the Child or Children
// functions) the WithChild function must be used.
//
// If the parent already has a child of the given name this function panics.
func NewChild(parent context.Context, name string) context.Context {
if Child(parent, name) != nil {
panic(fmt.Sprintf("child with name %q already exists on parent", name))
}
childPath := append(pathCP(parent), name)
child := withoutLocalValues(parent)
child = context.WithValue(child, ancestryKeyChildren, nil) // unset children
child = context.WithValue(child, ancestryKeyChildrenMap, nil) // unset children
child = context.WithValue(child, ancestryKeyParent, parent)
child = context.WithValue(child, ancestryKeyPath, childPath)
return child
}
func isChild(parent, child context.Context) bool {
parentPath, childPath := Path(parent), Path(child)
if len(parentPath) != len(childPath)-1 {
return false
}
for i := range parentPath {
if parentPath[i] != childPath[i] {
return false
}
}
return true
}
// WithChild returns a modified parent which holds a reference to child in its
// Children list. If the child's name is already taken in the parent then this
// function panics.
func WithChild(parent, child context.Context) context.Context {
if !isChild(parent, child) {
panic(fmt.Sprintf("child cannot be kept by Context which is not its parent"))
}
name, _ := Name(child)
children, childrenMap := childrenCP(parent)
if _, ok := childrenMap[name]; ok {
panic(fmt.Sprintf("child with name %q already exists on parent", name))
}
children = append(children, child)
childrenMap[name] = len(children) - 1
parent = context.WithValue(parent, ancestryKeyChildren, children)
parent = context.WithValue(parent, ancestryKeyChildrenMap, childrenMap)
return parent
}
// BreadthFirstVisit visits this Context and all of its children, and their
// children, etc... in a breadth-first order. If the callback returns false then
// the function returns without visiting any more Contexts.
func BreadthFirstVisit(ctx context.Context, callback func(context.Context) bool) {
queue := []context.Context{ctx}
for len(queue) > 0 {
if !callback(queue[0]) {
return
}
for _, child := range Children(queue[0]) {
queue = append(queue, child)
}
queue = queue[1:]
}
}
////////////////////////////////////////////////////////////////////////////////
// local value stuff
type localValsKey int
type localVal struct {
prev *localVal
key, val interface{}
}
// WithLocalValue is like context.WithValue, but the stored value will not be
// present on any children created via NewChild. Local values must be retrieved
// with the LocalValue function in this package. Local values share a different
// namespace than the normal WithValue/Value values (i.e. they do not overlap).
func WithLocalValue(ctx context.Context, key, val interface{}) context.Context {
prev, _ := ctx.Value(localValsKey(0)).(*localVal)
return context.WithValue(ctx, localValsKey(0), &localVal{
prev: prev,
key: key, val: val,
})
}
func withoutLocalValues(ctx context.Context) context.Context {
return context.WithValue(ctx, localValsKey(0), nil)
}
// LocalValue returns the value for the given key which was set by a call to
// WithLocalValue, or nil if no value was set for the given key.
func LocalValue(ctx context.Context, key interface{}) interface{} {
lv, _ := ctx.Value(localValsKey(0)).(*localVal)
for {
if lv == nil {
return nil
} else if lv.key == key {
return lv.val
}
lv = lv.prev
}
}
func localValuesIter(ctx context.Context, callback func(key, val interface{})) {
m := map[interface{}]struct{}{}
lv, _ := ctx.Value(localValsKey(0)).(*localVal)
for {
if lv == nil {
return
} else if _, ok := m[lv.key]; !ok {
callback(lv.key, lv.val)
m[lv.key] = struct{}{}
}
lv = lv.prev
}
}
// LocalValues returns all key/value pairs which have been set on the Context
// via WithLocalValue.
func LocalValues(ctx context.Context) map[interface{}]interface{} {
m := map[interface{}]interface{}{}
localValuesIter(ctx, func(key, val interface{}) {
m[key] = val
})
return m
}

View File

@ -1,156 +0,0 @@
package mctx
import (
"context"
. "testing"
"github.com/mediocregopher/mediocre-go-lib/mtest/massert"
)
func TestInheritance(t *T) {
ctx := context.Background()
ctx1 := NewChild(ctx, "1")
ctx1a := NewChild(ctx1, "a")
ctx1b := NewChild(ctx1, "b")
ctx1 = WithChild(ctx1, ctx1a)
ctx1 = WithChild(ctx1, ctx1b)
ctx2 := NewChild(ctx, "2")
ctx = WithChild(ctx, ctx1)
ctx = WithChild(ctx, ctx2)
massert.Require(t,
massert.Length(Path(ctx), 0),
massert.Equal(Path(ctx1), []string{"1"}),
massert.Equal(Path(ctx1a), []string{"1", "a"}),
massert.Equal(Path(ctx1b), []string{"1", "b"}),
massert.Equal(Path(ctx2), []string{"2"}),
)
massert.Require(t,
massert.Equal([]context.Context{ctx1, ctx2}, Children(ctx)),
massert.Equal([]context.Context{ctx1a, ctx1b}, Children(ctx1)),
massert.Length(Children(ctx2), 0),
)
}
func TestBreadFirstVisit(t *T) {
ctx := context.Background()
ctx1 := NewChild(ctx, "1")
ctx1a := NewChild(ctx1, "a")
ctx1b := NewChild(ctx1, "b")
ctx1 = WithChild(ctx1, ctx1a)
ctx1 = WithChild(ctx1, ctx1b)
ctx2 := NewChild(ctx, "2")
ctx = WithChild(ctx, ctx1)
ctx = WithChild(ctx, ctx2)
{
got := make([]context.Context, 0, 5)
BreadthFirstVisit(ctx, func(ctx context.Context) bool {
got = append(got, ctx)
return true
})
massert.Require(t,
massert.Equal([]context.Context{ctx, ctx1, ctx2, ctx1a, ctx1b}, got),
)
}
{
got := make([]context.Context, 0, 3)
BreadthFirstVisit(ctx, func(ctx context.Context) bool {
if len(Path(ctx)) > 1 {
return false
}
got = append(got, ctx)
return true
})
massert.Require(t,
massert.Equal([]context.Context{ctx, ctx1, ctx2}, got),
)
}
}
func TestLocalValues(t *T) {
// test with no value set
ctx := context.Background()
massert.Require(t,
massert.Nil(LocalValue(ctx, "foo")),
massert.Length(LocalValues(ctx), 0),
)
// test basic value retrieval
ctx = WithLocalValue(ctx, "foo", "bar")
massert.Require(t,
massert.Equal("bar", LocalValue(ctx, "foo")),
massert.Equal(
map[interface{}]interface{}{"foo": "bar"},
LocalValues(ctx),
),
)
// test that doesn't conflict with WithValue
ctx = context.WithValue(ctx, "foo", "WithValue bar")
massert.Require(t,
massert.Equal("bar", LocalValue(ctx, "foo")),
massert.Equal("WithValue bar", ctx.Value("foo")),
massert.Equal(
map[interface{}]interface{}{"foo": "bar"},
LocalValues(ctx),
),
)
// test that child doesn't get values
child := NewChild(ctx, "child")
massert.Require(t,
massert.Equal("bar", LocalValue(ctx, "foo")),
massert.Nil(LocalValue(child, "foo")),
massert.Length(LocalValues(child), 0),
)
// test that values on child don't affect parent values
child = WithLocalValue(child, "foo", "child bar")
ctx = WithChild(ctx, child)
massert.Require(t,
massert.Equal("bar", LocalValue(ctx, "foo")),
massert.Equal("child bar", LocalValue(child, "foo")),
massert.Equal(
map[interface{}]interface{}{"foo": "bar"},
LocalValues(ctx),
),
massert.Equal(
map[interface{}]interface{}{"foo": "child bar"},
LocalValues(child),
),
)
// test that two With calls on the same context generate distinct contexts
childA := WithLocalValue(child, "foo2", "baz")
childB := WithLocalValue(child, "foo2", "buz")
massert.Require(t,
massert.Equal("bar", LocalValue(ctx, "foo")),
massert.Equal("child bar", LocalValue(child, "foo")),
massert.Nil(LocalValue(child, "foo2")),
massert.Equal("baz", LocalValue(childA, "foo2")),
massert.Equal("buz", LocalValue(childB, "foo2")),
massert.Equal(
map[interface{}]interface{}{"foo": "child bar", "foo2": "baz"},
LocalValues(childA),
),
massert.Equal(
map[interface{}]interface{}{"foo": "child bar", "foo2": "buz"},
LocalValues(childB),
),
)
// if a value overwrites a previous one the newer one should show in
// LocalValues
ctx = WithLocalValue(ctx, "foo", "barbar")
massert.Require(t,
massert.Equal("barbar", LocalValue(ctx, "foo")),
massert.Equal(
map[interface{}]interface{}{"foo": "barbar"},
LocalValues(ctx),
),
)
}

15
mctx/mctx.go Normal file
View File

@ -0,0 +1,15 @@
// Package mctx extends the builtin context package to add easy-to-use
// annotation functionality, which is useful for logging and errors.
//
// All functions and methods in this package are thread-safe unless otherwise
// noted.
//
// Annotations
//
// Annotations are a special case of key/values, where the data being
// stored is specifically runtime metadata which would be useful for logging,
// error output, etc... Annotation data might include an IP address of a
// connected client, a userID the client has authenticated as, the primary key
// of a row in a database being queried, etc...
//
package mctx