remove mutex lock in a goroutine lifecycle

context.go

@@ -5,6 +5,11 @@ import (
 	"sync"
 )
 
+const (
+	initialMaxGoroutineCount = 1024
+	extendUnit               = 128
+)
+
 var (
 	mgrRegistry    = make(map[*ContextManager]bool)
 	mgrRegistryMtx sync.RWMutex
@@ -20,21 +25,45 @@ type Values map[interface{}]interface{}
 // class of context variables. You should use NewContextManager for
 // construction.
 type ContextManager struct {
-	mtx    sync.Mutex
-	values map[uint]Values
+	extendLock               sync.RWMutex
+	extendUnit               uint32
+	values                   []Values
+	currentMaxGoroutineCount int
 }
 
-// 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)}
+type Option struct {
+	InitialMaxGoroutineCount int
+	ExtendUnit               int
+}
+
+func newContextManager(opt Option) *ContextManager {
+	mgr := &ContextManager{values: make([]Values, opt.InitialMaxGoroutineCount)}
+	mgr.currentMaxGoroutineCount = len(mgr.values)
+	mgr.extendUnit = uint32(opt.ExtendUnit)
 	mgrRegistryMtx.Lock()
 	defer mgrRegistryMtx.Unlock()
 	mgrRegistry[mgr] = true
 	return mgr
 }
 
+func NewContextManagerWithOption(opt Option) *ContextManager {
+	if opt.InitialMaxGoroutineCount == 0 {
+		opt.InitialMaxGoroutineCount = initialMaxGoroutineCount
+	}
+	if opt.ExtendUnit == 0 {
+		opt.ExtendUnit = extendUnit
+	}
+
+	return newContextManager(opt)
+}
+
+// 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 {
+	return newContextManager(Option{InitialMaxGoroutineCount: initialMaxGoroutineCount, ExtendUnit: extendUnit})
+}
+
 // 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.
@@ -60,14 +89,17 @@ func (m *ContextManager) SetValues(new_values Values, context_call func()) {
 	mutated_keys := make([]interface{}, 0, len(new_values))
 	mutated_vals := make(Values, len(new_values))
 
-	EnsureGoroutineId(func(gid uint) {
-		m.mtx.Lock()
-		state, found := m.values[gid]
-		if !found {
+	EnsureGoroutineId(func(gid uint32) {
+		var found bool
+		m.extendIfNeeded(gid)
+
+		state := m.values[gid]
+		if state != nil {
+			found = true
+		} else {
 			state = make(Values, len(new_values))
 			m.values[gid] = state
 		}
-		m.mtx.Unlock()
 
 		for key, new_val := range new_values {
 			mutated_keys = append(mutated_keys, key)
@@ -79,9 +111,7 @@ func (m *ContextManager) SetValues(new_values Values, context_call func()) {
 
 		defer func() {
 			if !found {
-				m.mtx.Lock()
-				delete(m.values, gid)
-				m.mtx.Unlock()
+				m.values[gid] = nil
 				return
 			}
 
@@ -108,11 +138,9 @@ func (m *ContextManager) GetValue(key interface{}) (
 		return nil, false
 	}
 
-	m.mtx.Lock()
-	state, found := m.values[gid]
-	m.mtx.Unlock()
+	state := m.values[gid]
 
-	if !found {
+	if state == nil {
 		return nil, false
 	}
 	value, ok = state[key]
@@ -124,9 +152,7 @@ func (m *ContextManager) getValues() Values {
 	if !ok {
 		return nil
 	}
-	m.mtx.Lock()
-	state, _ := m.values[gid]
-	m.mtx.Unlock()
+	state := m.values[gid]
 	return state
 }
 
@@ -151,3 +177,23 @@ func Go(cb func()) {
 
 	go cb()
 }
+
+func (m *ContextManager) extend(gid uint32) {
+	m.extendLock.Lock()
+	defer m.extendLock.Unlock()
+	if gid >= uint32(m.currentMaxGoroutineCount) {
+		unit := ((gid-uint32(m.currentMaxGoroutineCount))/extendUnit + 1) * extendUnit
+		m.values = append(m.values, make([]Values, unit)...)
+		m.currentMaxGoroutineCount += int(unit)
+	}
+}
+
+func (m *ContextManager) extendIfNeeded(gid uint32) {
+	m.extendLock.RLock()
+	if gid >= uint32(m.currentMaxGoroutineCount) {
+		m.extendLock.RUnlock()
+		m.extend(gid)
+	} else {
+		m.extendLock.RUnlock()
+	}
+}

context_test.go

@@ -124,22 +124,38 @@ func ExampleGo() {
 
 func BenchmarkGetValue(b *testing.B) {
 	mgr := gls.NewContextManager()
+	wg := sync.WaitGroup{}
 	mgr.SetValues(gls.Values{"test_key": "test_val"}, func() {
 		b.ResetTimer()
 		for i := 0; i < b.N; i++ {
-			val, ok := mgr.GetValue("test_key")
-			if !ok || val != "test_val" {
-				b.FailNow()
+			for j := 0; j < 100; j++ {
+				wg.Add(1)
+				gls.Go(func() {
+					defer wg.Done()
+					val, ok := mgr.GetValue("test_key")
+					if !ok || val != "test_val" {
+						b.FailNow()
+					}
+				})
 			}
+			wg.Wait()
 		}
 	})
 }
 
 func BenchmarkSetValues(b *testing.B) {
 	mgr := gls.NewContextManager()
+	wg := sync.WaitGroup{}
 	for i := 0; i < b.N/2; i++ {
-		mgr.SetValues(gls.Values{"test_key": "test_val"}, func() {
-			mgr.SetValues(gls.Values{"test_key2": "test_val2"}, func() {})
-		})
+		for j := 0; j < 100; j++ {
+			wg.Add(1)
+			go func() {
+				defer wg.Done()
+				mgr.SetValues(gls.Values{"test_key": "test_val"}, func() {
+					mgr.SetValues(gls.Values{"test_key2": "test_val2"}, func() {})
+				})
+			}()
+		}
+		wg.Wait()
 	}
 }

gid.go

@@ -4,17 +4,23 @@ var (
 	stackTagPool = &idPool{}
 )
 
+func initIdPool() {
+	stackTagPool.Pool.New = func() interface{} {
+		return stackTagPool.newID()
+	}
+}
+
 // 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) {
+func GetGoroutineId() (gid uint32, 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)) {
+func EnsureGoroutineId(cb func(gid uint32)) {
 	if gid, ok := readStackTag(); ok {
 		cb(gid)
 		return

id_pool.go

@@ -6,29 +6,23 @@ package gls
 
 import (
 	"sync"
+	"sync/atomic"
 )
 
 type idPool struct {
-	mtx      sync.Mutex
-	released []uint
-	max_id   uint
+	sync.Pool
+	curID uint32
 }
 
-func (p *idPool) Acquire() (id uint) {
-	p.mtx.Lock()
-	defer p.mtx.Unlock()
-	if len(p.released) > 0 {
-		id = p.released[len(p.released)-1]
-		p.released = p.released[:len(p.released)-1]
-		return id
-	}
-	id = p.max_id
-	p.max_id++
-	return id
+func (p *idPool) newID() uint32 {
+	curID := atomic.AddUint32(&p.curID, 1)
+	return curID - 1
 }
 
-func (p *idPool) Release(id uint) {
-	p.mtx.Lock()
-	defer p.mtx.Unlock()
-	p.released = append(p.released, id)
+func (p *idPool) Acquire() (id uint32) {
+	return p.Get().(uint32)
+}
+
+func (p *idPool) Release(id uint32) {
+	p.Put(id)
 }

stack_tags.go

@@ -9,11 +9,11 @@ const (
 
 var (
 	pc_lookup   = make(map[uintptr]int8, 17)
-	mark_lookup [16]func(uint, func())
+	mark_lookup [16]func(uint32, func())
 )
 
 func init() {
-	setEntries := func(f func(uint, func()), v int8) {
+	setEntries := func(f func(uint32, func()), v int8) {
 		var ptr uintptr
 		f(0, func() {
 			ptr = findPtr()
@@ -40,9 +40,11 @@ func init() {
 	setEntries(github_com_jtolds_gls_markD, 0xd)
 	setEntries(github_com_jtolds_gls_markE, 0xe)
 	setEntries(github_com_jtolds_gls_markF, 0xf)
+
+	initIdPool()
 }
 
-func addStackTag(tag uint, context_call func()) {
+func addStackTag(tag uint32, context_call func()) {
 	if context_call == nil {
 		return
 	}
@@ -53,66 +55,66 @@ func addStackTag(tag uint, context_call func()) {
 // is easier. it shouldn't add any runtime cost in non-js builds.
 
 //go:noinline
-func github_com_jtolds_gls_markS(tag uint, cb func()) { _m(tag, cb) }
+func github_com_jtolds_gls_markS(tag uint32, cb func()) { _m(tag, cb) }
 
 //go:noinline
-func github_com_jtolds_gls_mark0(tag uint, cb func()) { _m(tag, cb) }
+func github_com_jtolds_gls_mark0(tag uint32, cb func()) { _m(tag, cb) }
 
 //go:noinline
-func github_com_jtolds_gls_mark1(tag uint, cb func()) { _m(tag, cb) }
+func github_com_jtolds_gls_mark1(tag uint32, cb func()) { _m(tag, cb) }
 
 //go:noinline
-func github_com_jtolds_gls_mark2(tag uint, cb func()) { _m(tag, cb) }
+func github_com_jtolds_gls_mark2(tag uint32, cb func()) { _m(tag, cb) }
 
 //go:noinline
-func github_com_jtolds_gls_mark3(tag uint, cb func()) { _m(tag, cb) }
+func github_com_jtolds_gls_mark3(tag uint32, cb func()) { _m(tag, cb) }
 
 //go:noinline
-func github_com_jtolds_gls_mark4(tag uint, cb func()) { _m(tag, cb) }
+func github_com_jtolds_gls_mark4(tag uint32, cb func()) { _m(tag, cb) }
 
 //go:noinline
-func github_com_jtolds_gls_mark5(tag uint, cb func()) { _m(tag, cb) }
+func github_com_jtolds_gls_mark5(tag uint32, cb func()) { _m(tag, cb) }
 
 //go:noinline
-func github_com_jtolds_gls_mark6(tag uint, cb func()) { _m(tag, cb) }
+func github_com_jtolds_gls_mark6(tag uint32, cb func()) { _m(tag, cb) }
 
 //go:noinline
-func github_com_jtolds_gls_mark7(tag uint, cb func()) { _m(tag, cb) }
+func github_com_jtolds_gls_mark7(tag uint32, cb func()) { _m(tag, cb) }
 
 //go:noinline
-func github_com_jtolds_gls_mark8(tag uint, cb func()) { _m(tag, cb) }
+func github_com_jtolds_gls_mark8(tag uint32, cb func()) { _m(tag, cb) }
 
 //go:noinline
-func github_com_jtolds_gls_mark9(tag uint, cb func()) { _m(tag, cb) }
+func github_com_jtolds_gls_mark9(tag uint32, cb func()) { _m(tag, cb) }
 
 //go:noinline
-func github_com_jtolds_gls_markA(tag uint, cb func()) { _m(tag, cb) }
+func github_com_jtolds_gls_markA(tag uint32, cb func()) { _m(tag, cb) }
 
 //go:noinline
-func github_com_jtolds_gls_markB(tag uint, cb func()) { _m(tag, cb) }
+func github_com_jtolds_gls_markB(tag uint32, cb func()) { _m(tag, cb) }
 
 //go:noinline
-func github_com_jtolds_gls_markC(tag uint, cb func()) { _m(tag, cb) }
+func github_com_jtolds_gls_markC(tag uint32, cb func()) { _m(tag, cb) }
 
 //go:noinline
-func github_com_jtolds_gls_markD(tag uint, cb func()) { _m(tag, cb) }
+func github_com_jtolds_gls_markD(tag uint32, cb func()) { _m(tag, cb) }
 
 //go:noinline
-func github_com_jtolds_gls_markE(tag uint, cb func()) { _m(tag, cb) }
+func github_com_jtolds_gls_markE(tag uint32, cb func()) { _m(tag, cb) }
 
 //go:noinline
-func github_com_jtolds_gls_markF(tag uint, cb func()) { _m(tag, cb) }
+func github_com_jtolds_gls_markF(tag uint32, cb func()) { _m(tag, cb) }
 
-func _m(tag_remainder uint, cb func()) {
+func _m(tag_remainder uint32, cb func()) {
 	if tag_remainder == 0 {
 		cb()
 	} else {
 		mark_lookup[tag_remainder&0xf](tag_remainder>>bitWidth, cb)
 	}
 }
 
-func readStackTag() (tag uint, ok bool) {
-	var current_tag uint
+func readStackTag() (tag uint32, ok bool) {
+	var current_tag uint32
 	offset := 0
 	for {
 		batch, next_offset := getStack(offset, stackBatchSize)
@@ -125,7 +127,7 @@ func readStackTag() (tag uint, ok bool) {
 				return current_tag, true
 			}
 			current_tag <<= bitWidth
-			current_tag += uint(val)
+			current_tag += uint32(val)
 		}
 		if next_offset == 0 {
 			break