github.com/hsfzxjy/dgo/go@v0.2.0/port_map.go (about)

     1  package dgo
     2  
     3  import (
     4  	"fmt"
     5  
     6  	xsync "github.com/puzpuzpuz/xsync/v2"
     7  )
     8  
     9  type _PortMap struct {
    10  	key2port    map[PortKey]*Port
    11  	defaultPort *Port
    12  	mu          *xsync.RBMutex
    13  }
    14  
    15  var portMap = newPortMap()
    16  
    17  func newPortMap() *_PortMap {
    18  	return &_PortMap{
    19  		key2port: make(map[PortKey]*Port),
    20  		mu:       xsync.NewRBMutex(),
    21  	}
    22  }
    23  
    24  func (m *_PortMap) addPort(sendPortKey PortKey, isDefault bool) {
    25  	m.mu.Lock()
    26  	defer m.mu.Unlock()
    27  
    28  	if isDefault && m.defaultPort != nil {
    29  		panic("dgo:go: default port is already set")
    30  	}
    31  
    32  	receivePortKey := m.initPort(sendPortKey)
    33  
    34  	port := newPort(
    35  		PortKey(sendPortKey),
    36  		PortKey(receivePortKey),
    37  	)
    38  	m.key2port[port.key()] = port
    39  
    40  	if isDefault {
    41  		m.defaultPort = port
    42  	}
    43  }
    44  
    45  func (m *_PortMap) removePort(port *Port) {
    46  	m.mu.Lock()
    47  	defer m.mu.Unlock()
    48  
    49  	key := port.key()
    50  	delete(m.key2port, key)
    51  	if m.defaultPort == port {
    52  		m.defaultPort = nil
    53  	}
    54  }
    55  
    56  func (m *_PortMap) DefaultPort() *Port {
    57  	rtok := m.mu.RLock()
    58  	defer m.mu.RUnlock(rtok)
    59  	return m.defaultPort
    60  }
    61  
    62  func (m *_PortMap) GetByPortKey(key PortKey) *Port {
    63  	rtok := m.mu.RLock()
    64  	defer m.mu.RUnlock(rtok)
    65  	if port, ok := m.key2port[key]; ok {
    66  		return port
    67  	} else {
    68  		panic(fmt.Sprintf("dgo:go: port [key=%d] not exists", key))
    69  	}
    70  }
    71  
    72  func (m *_PortMap) ResolvePort(port *Port) *Port {
    73  	if port != nil {
    74  		return port
    75  	}
    76  	return m.DefaultPort()
    77  }