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

     1  package client
     2  
     3  import (
     4  	"errors"
     5  	"fmt"
     6  	"github.com/nyan233/littlerpc/core/client/loadbalance"
     7  	"github.com/nyan233/littlerpc/core/common/msgparser"
     8  	"github.com/nyan233/littlerpc/core/common/transport"
     9  	"github.com/nyan233/littlerpc/core/utils/random"
    10  	"net"
    11  	"sync"
    12  	"sync/atomic"
    13  )
    14  
    15  // 严格来说, 这个对象不应该被释放, 它所使用的资源都应该是可被重置的
    16  // 当conn被关闭时, 新的conn可以复用被关闭的conn, 这个对象应该直到客户端被关闭之前
    17  // 一直存在于池中
    18  type connSource struct {
    19  	// 内嵌Conn类型使得可以作为负载均衡器的Value, 可减少一次并发map查找
    20  	transport.ConnAdapter
    21  	// 该flag被设置为true时表示该连接已经从负载均衡器剔除, 之后的连接已经使用不到了
    22  	// 但是可能有连接还在使用, 所以需要OnMsg过程在处理完该连接上的所有等待者需要的报文时关闭连接
    23  	// 并且阻止新的通知通道的绑定, 但是不阻止新消息的写入
    24  	halfClosed atomic.Bool
    25  	localAddr  net.Addr
    26  	remoteAddr net.Addr
    27  	// 表示该连接属于哪个节点
    28  	node loadbalance.RpcNode
    29  	// message ID的起始, 开始时随机分配
    30  	initSeq uint64
    31  	// 负责消息的解析
    32  	parser msgparser.Parser
    33  	mu     sync.Mutex
    34  	// 用于事件循环读取完毕的通知
    35  	notifySet map[uint64]chan Complete
    36  }
    37  
    38  func newConnSource(msgFactory msgparser.Factory, conn transport.ConnAdapter, node loadbalance.RpcNode) *connSource {
    39  	return &connSource{
    40  		ConnAdapter: conn,
    41  		localAddr:   conn.LocalAddr(),
    42  		remoteAddr:  conn.RemoteAddr(),
    43  		node:        node,
    44  		parser:      msgFactory(&msgparser.SimpleAllocTor{SharedPool: sharedPool.TakeMessagePool()}, 4096),
    45  		initSeq:     uint64(random.FastRand()),
    46  		notifySet:   make(map[uint64]chan Complete, 1024),
    47  	}
    48  }
    49  
    50  func (lc *connSource) GetMsgId() uint64 {
    51  	return atomic.AddUint64(&lc.initSeq, 1)
    52  }
    53  
    54  func (lc *connSource) SwapNotifyChannel(notify map[uint64]chan Complete) map[uint64]chan Complete {
    55  	lc.mu.Lock()
    56  	old := lc.notifySet
    57  	lc.notifySet = notify
    58  	lc.mu.Unlock()
    59  	return old
    60  }
    61  
    62  func (lc *connSource) BindNotifyChannel(msgId uint64, channel chan Complete) bool {
    63  	lc.mu.Lock()
    64  	if lc.notifySet == nil || lc.isHalfClosed() {
    65  		return false
    66  	}
    67  	lc.notifySet[msgId] = channel
    68  	lc.mu.Unlock()
    69  	return true
    70  }
    71  
    72  func (lc *connSource) PushCompleteMessage(msgId uint64, msg Complete) (end bool, err error) {
    73  	lc.mu.Lock()
    74  	defer lc.mu.Unlock()
    75  	if lc.notifySet == nil {
    76  		return true, errors.New("LRPC: connection already closed, notify set not found")
    77  	}
    78  	channel, ok := lc.notifySet[msgId]
    79  	if !ok {
    80  		return false, fmt.Errorf("LRPC: msgid is not found, msgid = %d", msgId)
    81  	}
    82  	select {
    83  	case channel <- msg:
    84  		break
    85  	default:
    86  		return false, errors.New("LRPC: notify channel is block state")
    87  	}
    88  	lc.notifySet[msgId] = nil
    89  	delete(lc.notifySet, msgId)
    90  	return len(lc.notifySet) == 0, nil
    91  }
    92  
    93  func (lc *connSource) HaveNWait() int {
    94  	lc.mu.Lock()
    95  	defer lc.mu.Unlock()
    96  	return len(lc.notifySet)
    97  }
    98  
    99  func (lc *connSource) halfClose() (ableRelease bool, err error) {
   100  	if !lc.halfClosed.CompareAndSwap(false, true) {
   101  		return false, errors.New("LRPC: connection already closed")
   102  	}
   103  	return lc.HaveNWait() <= 0, nil
   104  }
   105  
   106  func (lc *connSource) isHalfClosed() bool {
   107  	return lc.halfClosed.Load()
   108  }