context.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 globalc context.
	PutGlobal(key string, value interface{})

	// PutGlobalDynamic puts a key->value pair into the global context wwhere the value is
	// generated by a function that gets evaluated at every Read.
	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
	global   Map
	allmx    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.
	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 {
	id := curGoroutineID()
	cm.allmx.Lock()
	c := cm.contexts[id]
	if c == nil {
		c = cm.makeContext(id, nil, nil)
		cm.contexts[id] = c
		cm.allmx.Unlock()
		return c
	}
	cm.allmx.Unlock()
	return c.Enter()
}

func (c *context) Enter() Context {
	c.mx.RLock()
	id := c.id
	c.mx.RUnlock()
	next := c.cm.makeContext(id, c, nil)
	c.cm.allmx.Lock()
	c.cm.contexts[id] = next
	c.cm.allmx.Unlock()
	return next
}

func (c *context) Go(fn func()) {
	go func() {
		id := curGoroutineID()
		next := c.cm.makeContext(id, nil, c)
		c.cm.allmx.Lock()
		c.cm.contexts[id] = next
		c.cm.allmx.Unlock()
		fn()
		// Clean up the context
		c.cm.allmx.Lock()
		delete(c.cm.contexts, id)
		c.cm.allmx.Unlock()
	}()
}

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()
	if parent == nil {
		c.cm.allmx.Lock()
		delete(c.cm.contexts, id)
		c.cm.allmx.Unlock()
		return
	}
	c.cm.allmx.Lock()
	c.cm.contexts[id] = parent
	c.cm.allmx.Unlock()
}

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 (cm *manager) PutGlobal(key string, value interface{}) {
	cm.allmx.Lock()
	cm.global[key] = value
	cm.allmx.Unlock()
}

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) PutGlobalDynamic(key string, valueFN func() interface{}) {
	value := &dynval{valueFN}
	cm.allmx.Lock()
	cm.global[key] = value
	cm.allmx.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.allmx.RLock()
		fill(result, cm.global)
		cm.allmx.RUnlock()
	}
	return result
}

func fill(m Map, from Map) {
	if m != nil {
		for key, value := range from {
			_, alreadyRead := m[key]
			if !alreadyRead {
				switch v := value.(type) {
				case *dynval:
					m[key] = v.fn()
				default:
					m[key] = v
				}
			}
		}
	}
}

func (cm *manager) currentContext() *context {
	id := curGoroutineID()
	cm.allmx.RLock()
	c := cm.contexts[id]
	cm.allmx.RUnlock()
	return c
}