310 lines
7.0 KiB
Go
310 lines
7.0 KiB
Go
// Package context provides a mechanism for transparently tracking contextual
|
|
// state associated to the current goroutine and even across goroutines.
|
|
package context
|
|
|
|
import (
|
|
"sync"
|
|
)
|
|
|
|
// Manager provides the ability to create and access Contexts.
|
|
type Manager interface {
|
|
// Enter enters a new level on the current Context stack, creating a new Context
|
|
// if necessary.
|
|
Enter() Context
|
|
|
|
// Go starts the given function on a new goroutine but sharing the context of
|
|
// the current goroutine (if it has one).
|
|
Go(func())
|
|
|
|
// PutGlobal puts the given key->value pair into the global context.
|
|
PutGlobal(key string, value interface{})
|
|
|
|
// PutGlobalDynamic puts a key->value pair into the global context where the
|
|
// value is generated by a function that gets evaluated at every Read. If the
|
|
// value is a map[string]interface{}, we will unpack the map and set each
|
|
// contained key->value pair independently.
|
|
PutGlobalDynamic(key string, valueFN func() interface{})
|
|
|
|
// AsMap returns a map containing all values from the supplied obj if it is a
|
|
// Contextual, plus any addition values from along the stack, plus globals if so
|
|
// specified.
|
|
AsMap(obj interface{}, includeGlobals bool) Map
|
|
}
|
|
|
|
type manager struct {
|
|
contexts map[uint64]*context
|
|
mxContexts sync.RWMutex
|
|
global Map
|
|
mxGlobal sync.RWMutex
|
|
}
|
|
|
|
// NewManager creates a new Manager
|
|
func NewManager() Manager {
|
|
return &manager{
|
|
contexts: make(map[uint64]*context),
|
|
global: make(Map),
|
|
}
|
|
}
|
|
|
|
// Contextual is an interface for anything that maintains its own context.
|
|
type Contextual interface {
|
|
// Fill fills the given Map with all of this Contextual's context
|
|
Fill(m Map)
|
|
}
|
|
|
|
// Map is a map of key->value pairs.
|
|
type Map map[string]interface{}
|
|
|
|
// Fill implements the method from the Contextual interface.
|
|
func (_m Map) Fill(m Map) {
|
|
for key, value := range _m {
|
|
m[key] = value
|
|
}
|
|
}
|
|
|
|
// Context is a context containing key->value pairs
|
|
type Context interface {
|
|
// Enter enters a new level on this Context stack.
|
|
Enter() Context
|
|
|
|
// Go starts the given function on a new goroutine.
|
|
Go(fn func())
|
|
|
|
// Exit exits the current level on this Context stack.
|
|
Exit()
|
|
|
|
// Put puts a key->value pair into the current level of the context stack.
|
|
Put(key string, value interface{}) Context
|
|
|
|
// PutIfAbsent puts the given key->value pair into the current level of the
|
|
// context stack if and only if that key is defined nowhere within the context
|
|
// stack (including parent contexts).
|
|
PutIfAbsent(key string, value interface{}) Context
|
|
|
|
// PutDynamic puts a key->value pair into the current level of the context
|
|
// stack where the value is generated by a function that gets evaluated at
|
|
// every Read. If the value is a map[string]interface{}, we will unpack the
|
|
// map and set each contained key->value pair independently.
|
|
PutDynamic(key string, valueFN func() interface{}) Context
|
|
|
|
// Fill fills the given map with data from this Context
|
|
Fill(m Map)
|
|
|
|
// AsMap returns a map containing all values from the supplied obj if it is a
|
|
// Contextual, plus any addition values from along the stack, plus globals if
|
|
// so specified.
|
|
AsMap(obj interface{}, includeGlobals bool) Map
|
|
}
|
|
|
|
type context struct {
|
|
cm *manager
|
|
id uint64
|
|
parent *context
|
|
branchedFrom *context
|
|
data Map
|
|
mx sync.RWMutex
|
|
}
|
|
|
|
type dynval struct {
|
|
fn func() interface{}
|
|
}
|
|
|
|
func (cm *manager) Enter() Context {
|
|
return cm.enter(curGoroutineID())
|
|
}
|
|
|
|
func (cm *manager) enter(id uint64) *context {
|
|
cm.mxContexts.Lock()
|
|
parentOrNil := cm.contexts[id]
|
|
c := cm.makeContext(id, parentOrNil, nil)
|
|
cm.contexts[id] = c
|
|
cm.mxContexts.Unlock()
|
|
return c
|
|
}
|
|
|
|
func (cm *manager) exit(id uint64, parent *context) {
|
|
cm.mxContexts.Lock()
|
|
if parent == nil {
|
|
delete(cm.contexts, id)
|
|
} else {
|
|
cm.contexts[id] = parent
|
|
}
|
|
cm.mxContexts.Unlock()
|
|
}
|
|
|
|
func (cm *manager) branch(id uint64, from *context) {
|
|
next := cm.makeContext(id, nil, from)
|
|
cm.mxContexts.Lock()
|
|
cm.contexts[id] = next
|
|
cm.mxContexts.Unlock()
|
|
}
|
|
|
|
func (cm *manager) merge(id uint64) {
|
|
cm.mxContexts.Lock()
|
|
delete(cm.contexts, id)
|
|
cm.mxContexts.Unlock()
|
|
}
|
|
|
|
func (c *context) Enter() Context {
|
|
c.mx.RLock()
|
|
id := c.id
|
|
c.mx.RUnlock()
|
|
return c.cm.enter(id)
|
|
}
|
|
|
|
func (c *context) Go(fn func()) {
|
|
go func() {
|
|
id := curGoroutineID()
|
|
c.cm.branch(id, c)
|
|
fn()
|
|
c.cm.merge(id)
|
|
}()
|
|
}
|
|
|
|
func (cm *manager) Go(fn func()) {
|
|
c := cm.currentContext()
|
|
if c != nil {
|
|
c.Go(fn)
|
|
} else {
|
|
go fn()
|
|
}
|
|
}
|
|
|
|
func (cm *manager) makeContext(id uint64, parent *context, branchedFrom *context) *context {
|
|
return &context{
|
|
cm: cm,
|
|
id: id,
|
|
parent: parent,
|
|
branchedFrom: branchedFrom,
|
|
data: make(Map),
|
|
}
|
|
}
|
|
|
|
func (c *context) Exit() {
|
|
c.mx.RLock()
|
|
id := c.id
|
|
parent := c.parent
|
|
c.mx.RUnlock()
|
|
c.cm.exit(id, parent)
|
|
}
|
|
|
|
func (c *context) Put(key string, value interface{}) Context {
|
|
c.mx.Lock()
|
|
c.data[key] = value
|
|
c.mx.Unlock()
|
|
return c
|
|
}
|
|
|
|
func (c *context) PutIfAbsent(key string, value interface{}) Context {
|
|
for ctx := c; ctx != nil; {
|
|
ctx.mx.RLock()
|
|
_, exists := ctx.data[key]
|
|
next := ctx.parent
|
|
if next == nil {
|
|
next = ctx.branchedFrom
|
|
}
|
|
ctx.mx.RUnlock()
|
|
if exists {
|
|
return c
|
|
}
|
|
ctx = next
|
|
}
|
|
|
|
// Value not set, set it
|
|
return c.Put(key, value)
|
|
}
|
|
|
|
func (c *context) PutDynamic(key string, valueFN func() interface{}) Context {
|
|
value := &dynval{valueFN}
|
|
c.mx.Lock()
|
|
c.data[key] = value
|
|
c.mx.Unlock()
|
|
return c
|
|
}
|
|
|
|
func (cm *manager) PutGlobal(key string, value interface{}) {
|
|
cm.mxGlobal.Lock()
|
|
cm.global[key] = value
|
|
cm.mxGlobal.Unlock()
|
|
}
|
|
|
|
func (cm *manager) PutGlobalDynamic(key string, valueFN func() interface{}) {
|
|
value := &dynval{valueFN}
|
|
cm.mxGlobal.Lock()
|
|
cm.global[key] = value
|
|
cm.mxGlobal.Unlock()
|
|
}
|
|
|
|
func (c *context) Fill(m Map) {
|
|
for ctx := c; ctx != nil; {
|
|
ctx.mx.RLock()
|
|
fill(m, ctx.data)
|
|
next := ctx.parent
|
|
if next == nil {
|
|
next = ctx.branchedFrom
|
|
}
|
|
ctx.mx.RUnlock()
|
|
ctx = next
|
|
}
|
|
}
|
|
|
|
func (cm *manager) AsMap(obj interface{}, includeGlobals bool) Map {
|
|
return cm.currentContext().asMap(cm, obj, includeGlobals)
|
|
}
|
|
|
|
func (c *context) AsMap(obj interface{}, includeGlobals bool) Map {
|
|
return c.asMap(c.cm, obj, includeGlobals)
|
|
}
|
|
|
|
func (c *context) asMap(cm *manager, obj interface{}, includeGlobals bool) Map {
|
|
result := make(Map, 0)
|
|
cl, ok := obj.(Contextual)
|
|
if ok {
|
|
cl.Fill(result)
|
|
}
|
|
if c != nil {
|
|
c.Fill(result)
|
|
}
|
|
if includeGlobals {
|
|
cm.mxGlobal.RLock()
|
|
fill(result, cm.global)
|
|
cm.mxGlobal.RUnlock()
|
|
}
|
|
return result
|
|
}
|
|
|
|
func fill(m Map, from Map) {
|
|
if m != nil {
|
|
doFill := func(key string, _value interface{}) {
|
|
switch value := _value.(type) {
|
|
case map[string]interface{}:
|
|
for k, v := range value {
|
|
m[k] = v
|
|
}
|
|
default:
|
|
m[key] = value
|
|
}
|
|
}
|
|
|
|
for key, value := range from {
|
|
_, alreadyRead := m[key]
|
|
if !alreadyRead {
|
|
switch v := value.(type) {
|
|
case *dynval:
|
|
doFill(key, v.fn())
|
|
default:
|
|
doFill(key, v)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func (cm *manager) currentContext() *context {
|
|
id := curGoroutineID()
|
|
cm.mxContexts.RLock()
|
|
c := cm.contexts[id]
|
|
cm.mxContexts.RUnlock()
|
|
return c
|
|
}
|