|
|
@@ -0,0 +1,150 @@
|
|
|
+// Package gls implements goroutine-local storage.
|
|
|
+package gls
|
|
|
+
|
|
|
+import (
|
|
|
+ "runtime"
|
|
|
+ "sync"
|
|
|
+)
|
|
|
+
|
|
|
+const (
|
|
|
+ maxCallers = 64
|
|
|
+)
|
|
|
+
|
|
|
+var (
|
|
|
+ stackTagPool = &idPool{}
|
|
|
+ mgrRegistry = make(map[*ContextManager]bool)
|
|
|
+ mgrRegistryMtx sync.RWMutex
|
|
|
+)
|
|
|
+
|
|
|
+// Values is simply a map of key types to value types. Used by SetValues to
|
|
|
+// set multiple values at once.
|
|
|
+type Values map[interface{}]interface{}
|
|
|
+
|
|
|
+func currentStack(skip int) []uintptr {
|
|
|
+ stack := make([]uintptr, maxCallers)
|
|
|
+ return stack[:runtime.Callers(2+skip, stack)]
|
|
|
+}
|
|
|
+
|
|
|
+// ContextManager is the main entrypoint for interacting with
|
|
|
+// Goroutine-local-storage. You can have multiple independent ContextManagers
|
|
|
+// at any given time. ContextManagers are usually declared globally for a given
|
|
|
+// class of context variables. You should use NewContextManager for
|
|
|
+// construction.
|
|
|
+type ContextManager struct {
|
|
|
+ mtx sync.RWMutex
|
|
|
+ values map[uint]Values
|
|
|
+}
|
|
|
+
|
|
|
+// NewContextManager returns a brand new ContextManager. It also registers the
|
|
|
+// new ContextManager in the ContextManager registry which is used by the Go
|
|
|
+// method. ContextManagers are typically defined globally at package scope.
|
|
|
+func NewContextManager() *ContextManager {
|
|
|
+ mgr := &ContextManager{values: make(map[uint]Values)}
|
|
|
+ mgrRegistryMtx.Lock()
|
|
|
+ defer mgrRegistryMtx.Unlock()
|
|
|
+ mgrRegistry[mgr] = true
|
|
|
+ return mgr
|
|
|
+}
|
|
|
+
|
|
|
+// Unregister removes a ContextManager from the global registry, used by the
|
|
|
+// Go method. Only intended for use when you're completely done with a
|
|
|
+// ContextManager. Use of Unregister at all is rare.
|
|
|
+func (m *ContextManager) Unregister() {
|
|
|
+ mgrRegistryMtx.Lock()
|
|
|
+ defer mgrRegistryMtx.Unlock()
|
|
|
+ delete(mgrRegistry, m)
|
|
|
+}
|
|
|
+
|
|
|
+// SetValues takes a collection of values and a function to call for those
|
|
|
+// values to be set in. Anything further down the stack will have the set
|
|
|
+// values available through GetValue. SetValues will add new values or replace
|
|
|
+// existing values of the same key and will not mutate or change values for
|
|
|
+// previous stack frames.
|
|
|
+// SetValues is slow (makes a copy of all current and new values for the new
|
|
|
+// gls-context) in order to reduce the amount of lookups GetValue requires.
|
|
|
+func (m *ContextManager) SetValues(new_values Values, context_call func()) {
|
|
|
+ if len(new_values) == 0 {
|
|
|
+ context_call()
|
|
|
+ return
|
|
|
+ }
|
|
|
+
|
|
|
+ tags := readStackTags(currentStack(1))
|
|
|
+
|
|
|
+ m.mtx.Lock()
|
|
|
+ values := new_values
|
|
|
+ for _, tag := range tags {
|
|
|
+ if existing_values, ok := m.values[tag]; ok {
|
|
|
+ // oh, we found existing values, let's make a copy
|
|
|
+ values = make(Values, len(existing_values)+len(new_values))
|
|
|
+ for key, val := range existing_values {
|
|
|
+ values[key] = val
|
|
|
+ }
|
|
|
+ for key, val := range new_values {
|
|
|
+ values[key] = val
|
|
|
+ }
|
|
|
+ break
|
|
|
+ }
|
|
|
+ }
|
|
|
+ new_tag := stackTagPool.Acquire()
|
|
|
+ m.values[new_tag] = values
|
|
|
+ m.mtx.Unlock()
|
|
|
+ defer func() {
|
|
|
+ m.mtx.Lock()
|
|
|
+ delete(m.values, new_tag)
|
|
|
+ m.mtx.Unlock()
|
|
|
+ stackTagPool.Release(new_tag)
|
|
|
+ }()
|
|
|
+
|
|
|
+ addStackTag(new_tag, context_call)
|
|
|
+}
|
|
|
+
|
|
|
+// GetValue will return a previously set value, provided that the value was set
|
|
|
+// by SetValues somewhere higher up the stack. If the value is not found, ok
|
|
|
+// will be false.
|
|
|
+func (m *ContextManager) GetValue(key interface{}) (value interface{}, ok bool) {
|
|
|
+
|
|
|
+ tags := readStackTags(currentStack(1))
|
|
|
+ m.mtx.RLock()
|
|
|
+ defer m.mtx.RUnlock()
|
|
|
+ for _, tag := range tags {
|
|
|
+ if values, ok := m.values[tag]; ok {
|
|
|
+ value, ok := values[key]
|
|
|
+ return value, ok
|
|
|
+ }
|
|
|
+ }
|
|
|
+ return "", false
|
|
|
+}
|
|
|
+
|
|
|
+func (m *ContextManager) getValues() Values {
|
|
|
+ tags := readStackTags(currentStack(2))
|
|
|
+ m.mtx.RLock()
|
|
|
+ defer m.mtx.RUnlock()
|
|
|
+ for _, tag := range tags {
|
|
|
+ if values, ok := m.values[tag]; ok {
|
|
|
+ return values
|
|
|
+ }
|
|
|
+ }
|
|
|
+ return nil
|
|
|
+}
|
|
|
+
|
|
|
+// Go preserves ContextManager values and Goroutine-local-storage across new
|
|
|
+// goroutine invocations. The Go method makes a copy of all existing values on
|
|
|
+// all registered context managers and makes sure they are still set after
|
|
|
+// kicking off the provided function in a new goroutine. If you don't use this
|
|
|
+// Go method instead of the standard 'go' keyword, you will lose values in
|
|
|
+// ContextManagers, as goroutines have brand new stacks.
|
|
|
+func Go(cb func()) {
|
|
|
+ mgrRegistryMtx.RLock()
|
|
|
+ defer mgrRegistryMtx.RUnlock()
|
|
|
+
|
|
|
+ for mgr, _ := range mgrRegistry {
|
|
|
+ values := mgr.getValues()
|
|
|
+ if len(values) > 0 {
|
|
|
+ mgr_copy := mgr
|
|
|
+ cb_copy := cb
|
|
|
+ cb = func() { mgr_copy.SetValues(values, cb_copy) }
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ go cb()
|
|
|
+}
|