github.com/nyan233/littlerpc@v0.4.6-0.20230316182519-0c8d5c48abaf/core/client/context_manager.go (about)

     1  package client
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  )
     7  
     8  type contextManager struct {
     9  	lock     sync.Mutex
    10  	ctxIdSet map[context.Context]uint64
    11  }
    12  
    13  func newContextManager() *contextManager {
    14  	return &contextManager{
    15  		ctxIdSet: make(map[context.Context]uint64, 16),
    16  	}
    17  }
    18  
    19  // Register 返回有效的Id
    20  func (cm *contextManager) Register(ctx context.Context, id uint64) (_ uint64) {
    21  	cm.lock.Lock()
    22  	defer cm.lock.Unlock()
    23  	cid, ok := cm.ctxIdSet[ctx]
    24  	if ok {
    25  		return cid
    26  	}
    27  	cm.ctxIdSet[ctx] = id
    28  	return id
    29  }
    30  
    31  // Unregister 返回有效的Id
    32  func (cm *contextManager) Unregister(ctx context.Context) (_ uint64) {
    33  	cm.lock.Lock()
    34  	defer cm.lock.Unlock()
    35  	cid, ok := cm.ctxIdSet[ctx]
    36  	if !ok {
    37  		return 0
    38  	}
    39  	delete(cm.ctxIdSet, ctx)
    40  	return cid
    41  }