completely redo internals

a bunch of benefits:
 * no limit on stack size
 * now only adds stack tags once per goroutine
 * improved gopherjs support
 * expose GoroutineId values.

Author: jtolio

context.go

@@ -5,12 +5,7 @@ import (
 	"sync"
 )
 
-const (
-	maxCallers = 64
-)
-
 var (
-	stackTagPool   = &idPool{}
 	mgrRegistry    = make(map[*ContextManager]bool)
 	mgrRegistryMtx sync.RWMutex
 )
@@ -25,7 +20,7 @@ type Values map[interface{}]interface{}
 // class of context variables. You should use NewContextManager for
 // construction.
 type ContextManager struct {
-	mtx    sync.RWMutex
+	mtx    sync.Mutex
 	values map[uint]Values
 }
 
@@ -62,63 +57,77 @@ func (m *ContextManager) SetValues(new_values Values, context_call func()) {
 		return
 	}
 
-	tags := readStackTags(1)
+	mutated_keys := make([]interface{}, 0, len(new_values))
+	mutated_vals := make(Values, len(new_values))
 
-	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() {
+	EnsureGoroutineId(func(gid uint) {
 		m.mtx.Lock()
-		delete(m.values, new_tag)
+		state, found := m.values[gid]
+		if !found {
+			state = make(Values, len(new_values))
+			m.values[gid] = state
+		}
 		m.mtx.Unlock()
-		stackTagPool.Release(new_tag)
-	}()
 
-	addStackTag(new_tag, context_call)
+		for key, new_val := range new_values {
+			mutated_keys = append(mutated_keys, key)
+			if old_val, ok := state[key]; ok {
+				mutated_vals[key] = old_val
+			}
+			state[key] = new_val
+		}
+
+		defer func() {
+			if !found {
+				m.mtx.Lock()
+				delete(m.values, gid)
+				m.mtx.Unlock()
+				return
+			}
+
+			for _, key := range mutated_keys {
+				if val, ok := mutated_vals[key]; ok {
+					state[key] = val
+				} else {
+					delete(state, key)
+				}
+			}
+		}()
+
+		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(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
-		}
+func (m *ContextManager) GetValue(key interface{}) (
+	value interface{}, ok bool) {
+	gid, ok := GetGoroutineId()
+	if !ok {
+		return nil, false
 	}
-	return "", false
+
+	m.mtx.Lock()
+	state, found := m.values[gid]
+	m.mtx.Unlock()
+
+	if !found {
+		return nil, false
+	}
+	value, ok = state[key]
+	return value, ok
 }
 
 func (m *ContextManager) getValues() Values {
-	tags := readStackTags(2)
-	m.mtx.RLock()
-	defer m.mtx.RUnlock()
-	for _, tag := range tags {
-		if values, ok := m.values[tag]; ok {
-			return values
-		}
+	gid, ok := GetGoroutineId()
+	if !ok {
+		return nil
 	}
-	return nil
+	m.mtx.Lock()
+	state, _ := m.values[gid]
+	m.mtx.Unlock()
+	return state
 }
 
 // Go preserves ContextManager values and Goroutine-local-storage across new
@@ -131,12 +140,12 @@ func Go(cb func()) {
 	mgrRegistryMtx.RLock()
 	defer mgrRegistryMtx.RUnlock()
 
-	for mgr, _ := range mgrRegistry {
+	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) }
+			cb = func(mgr *ContextManager, cb func()) func() {
+				return func() { mgr.SetValues(values, cb) }
+			}(mgr, cb)
 		}
 	}
 

context_test.go

@@ -1,16 +1,18 @@
-package gls
+package gls_test
 
 import (
 	"fmt"
 	"sync"
 	"testing"
+
+	"github.com/jtolds/gls"
 )
 
 func TestContexts(t *testing.T) {
-	mgr1 := NewContextManager()
-	mgr2 := NewContextManager()
+	mgr1 := gls.NewContextManager()
+	mgr2 := gls.NewContextManager()
 
-	CheckVal := func(mgr *ContextManager, key, exp_val string) {
+	CheckVal := func(mgr *gls.ContextManager, key, exp_val string) {
 		val, ok := mgr.GetValue(key)
 		if len(exp_val) == 0 {
 			if ok {
@@ -37,32 +39,36 @@ func TestContexts(t *testing.T) {
 	}
 
 	Check("", "", "", "")
-	mgr2.SetValues(Values{"key1": "val1c"}, func() {
+	mgr2.SetValues(gls.Values{"key1": "val1c"}, func() {
 		Check("", "", "val1c", "")
-		mgr1.SetValues(Values{"key1": "val1a"}, func() {
+		mgr1.SetValues(gls.Values{"key1": "val1a"}, func() {
 			Check("val1a", "", "val1c", "")
-			mgr1.SetValues(Values{"key2": "val1b"}, func() {
+			mgr1.SetValues(gls.Values{"key2": "val1b"}, func() {
 				Check("val1a", "val1b", "val1c", "")
 				var wg sync.WaitGroup
 				wg.Add(2)
 				go func() {
 					defer wg.Done()
 					Check("", "", "", "")
 				}()
-				Go(func() {
+				gls.Go(func() {
 					defer wg.Done()
 					Check("val1a", "val1b", "val1c", "")
 				})
 				wg.Wait()
+				Check("val1a", "val1b", "val1c", "")
 			})
+			Check("val1a", "", "val1c", "")
 		})
+		Check("", "", "val1c", "")
 	})
+	Check("", "", "", "")
 }
 
 func ExampleContextManager_SetValues() {
 	var (
-		mgr            = NewContextManager()
-		request_id_key = GenSym()
+		mgr            = gls.NewContextManager()
+		request_id_key = gls.GenSym()
 	)
 
 	MyLog := func() {
@@ -73,7 +79,7 @@ func ExampleContextManager_SetValues() {
 		}
 	}
 
-	mgr.SetValues(Values{request_id_key: "12345"}, func() {
+	mgr.SetValues(gls.Values{request_id_key: "12345"}, func() {
 		MyLog()
 	})
 	MyLog()
@@ -84,8 +90,8 @@ func ExampleContextManager_SetValues() {
 
 func ExampleGo() {
 	var (
-		mgr            = NewContextManager()
-		request_id_key = GenSym()
+		mgr            = gls.NewContextManager()
+		request_id_key = gls.GenSym()
 	)
 
 	MyLog := func() {
@@ -96,7 +102,7 @@ func ExampleGo() {
 		}
 	}
 
-	mgr.SetValues(Values{request_id_key: "12345"}, func() {
+	mgr.SetValues(gls.Values{request_id_key: "12345"}, func() {
 		var wg sync.WaitGroup
 		wg.Add(1)
 		go func() {
@@ -105,7 +111,7 @@ func ExampleGo() {
 		}()
 		wg.Wait()
 		wg.Add(1)
-		Go(func() {
+		gls.Go(func() {
 			defer wg.Done()
 			MyLog()
 		})
@@ -117,8 +123,8 @@ func ExampleGo() {
 }
 
 func BenchmarkGetValue(b *testing.B) {
-	mgr := NewContextManager()
-	mgr.SetValues(Values{"test_key": "test_val"}, func() {
+	mgr := gls.NewContextManager()
+	mgr.SetValues(gls.Values{"test_key": "test_val"}, func() {
 		b.ResetTimer()
 		for i := 0; i < b.N; i++ {
 			val, ok := mgr.GetValue("test_key")
@@ -130,10 +136,10 @@ func BenchmarkGetValue(b *testing.B) {
 }
 
 func BenchmarkSetValues(b *testing.B) {
-	mgr := NewContextManager()
+	mgr := gls.NewContextManager()
 	for i := 0; i < b.N/2; i++ {
-		mgr.SetValues(Values{"test_key": "test_val"}, func() {
-			mgr.SetValues(Values{"test_key2": "test_val2"}, func() {})
+		mgr.SetValues(gls.Values{"test_key": "test_val"}, func() {
+			mgr.SetValues(gls.Values{"test_key2": "test_val2"}, func() {})
 		})
 	}
 }

gen_sym.go

@@ -1,13 +1,21 @@
 package gls
 
+import (
+	"sync"
+)
+
 var (
-	symPool = &idPool{}
+	keyMtx     sync.Mutex
+	keyCounter uint64
 )
 
 // ContextKey is a throwaway value you can use as a key to a ContextManager
-type ContextKey struct{ id uint }
+type ContextKey struct{ id uint64 }
 
 // GenSym will return a brand new, never-before-used ContextKey
 func GenSym() ContextKey {
-	return ContextKey{id: symPool.Acquire()}
+	keyMtx.Lock()
+	defer keyMtx.Unlock()
+	keyCounter += 1
+	return ContextKey{id: keyCounter}
 }

gid.go

@@ -0,0 +1,25 @@
+package gls
+
+var (
+	stackTagPool = &idPool{}
+)
+
+// Will return this goroutine's identifier if set. If you always need a
+// goroutine identifier, you should use EnsureGoroutineId which will make one
+// if there isn't one already.
+func GetGoroutineId() (gid uint, ok bool) {
+	return readStackTag()
+}
+
+// Will call cb with the current goroutine identifier. If one hasn't already
+// been generated, one will be created and set first. The goroutine identifier
+// might be invalid after cb returns.
+func EnsureGoroutineId(cb func(gid uint)) {
+	if gid, ok := readStackTag(); ok {
+		cb(gid)
+		return
+	}
+	gid := stackTagPool.Acquire()
+	defer stackTagPool.Release(gid)
+	addStackTag(gid, func() { cb(gid) })
+}