123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121 |
- package transport
- import (
- "context"
- "reflect"
- "sync"
- "github.com/fatedier/golib/errors"
- "github.com/fatedier/frp/pkg/msg"
- )
- type MessageTransporter interface {
- Send(msg.Message) error
-
-
- Do(ctx context.Context, req msg.Message, laneKey, recvMsgType string) (msg.Message, error)
-
- Dispatch(m msg.Message, laneKey string) bool
-
- DispatchWithType(m msg.Message, msgType, laneKey string) bool
- }
- func NewMessageTransporter(sendCh chan msg.Message) MessageTransporter {
- return &transporterImpl{
- sendCh: sendCh,
- registry: make(map[string]map[string]chan msg.Message),
- }
- }
- type transporterImpl struct {
- sendCh chan msg.Message
-
-
-
- registry map[string]map[string]chan msg.Message
- mu sync.RWMutex
- }
- func (impl *transporterImpl) Send(m msg.Message) error {
- return errors.PanicToError(func() {
- impl.sendCh <- m
- })
- }
- func (impl *transporterImpl) Do(ctx context.Context, req msg.Message, laneKey, recvMsgType string) (msg.Message, error) {
- ch := make(chan msg.Message, 1)
- defer close(ch)
- unregisterFn := impl.registerMsgChan(ch, laneKey, recvMsgType)
- defer unregisterFn()
- if err := impl.Send(req); err != nil {
- return nil, err
- }
- select {
- case <-ctx.Done():
- return nil, ctx.Err()
- case resp := <-ch:
- return resp, nil
- }
- }
- func (impl *transporterImpl) DispatchWithType(m msg.Message, msgType, laneKey string) bool {
- var ch chan msg.Message
- impl.mu.RLock()
- byLaneKey, ok := impl.registry[msgType]
- if ok {
- ch = byLaneKey[laneKey]
- }
- impl.mu.RUnlock()
- if ch == nil {
- return false
- }
- if err := errors.PanicToError(func() {
- ch <- m
- }); err != nil {
- return false
- }
- return true
- }
- func (impl *transporterImpl) Dispatch(m msg.Message, laneKey string) bool {
- msgType := reflect.TypeOf(m).Elem().Name()
- return impl.DispatchWithType(m, msgType, laneKey)
- }
- func (impl *transporterImpl) registerMsgChan(recvCh chan msg.Message, laneKey string, msgType string) (unregister func()) {
- impl.mu.Lock()
- byLaneKey, ok := impl.registry[msgType]
- if !ok {
- byLaneKey = make(map[string]chan msg.Message)
- impl.registry[msgType] = byLaneKey
- }
- byLaneKey[laneKey] = recvCh
- impl.mu.Unlock()
- unregister = func() {
- impl.mu.Lock()
- delete(byLaneKey, laneKey)
- impl.mu.Unlock()
- }
- return
- }
|