github.com/glide-im/glide@v1.6.0/pkg/gate/reader.go (about)

     1  package gate
     2  
     3  import (
     4  	"github.com/glide-im/glide/pkg/conn"
     5  	"github.com/glide-im/glide/pkg/logger"
     6  	"github.com/glide-im/glide/pkg/messages"
     7  	"sync"
     8  )
     9  
    10  var messageReader MessageReader
    11  
    12  // var codec messages.Codec = message.ProtobufCodec{}
    13  var codec messages.Codec = messages.DefaultCodec
    14  
    15  // recyclePool 回收池, 减少临时对象, 回收复用 readerRes
    16  var recyclePool sync.Pool
    17  
    18  func init() {
    19  	recyclePool = sync.Pool{
    20  		New: func() interface{} {
    21  			return &readerRes{}
    22  		},
    23  	}
    24  	SetMessageReader(&defaultReader{})
    25  }
    26  
    27  func SetMessageReader(s MessageReader) {
    28  	messageReader = s
    29  }
    30  
    31  type readerRes struct {
    32  	err error
    33  	m   *messages.GlideMessage
    34  }
    35  
    36  // Recycle 回收当前对象, 一定要在用完后调用这个方法, 否则无法回收
    37  func (r *readerRes) Recycle() {
    38  	r.m = nil
    39  	r.err = nil
    40  	recyclePool.Put(r)
    41  }
    42  
    43  // MessageReader 表示一个从连接中(Connection)读取消息的读取者, 可以用于定义如何从连接中读取并解析消息.
    44  type MessageReader interface {
    45  
    46  	// Read 阻塞读取, 会阻塞当前协程
    47  	Read(conn conn.Connection) (*messages.GlideMessage, error)
    48  
    49  	// ReadCh 返回两个管道, 第一个用于读取内容, 第二个用于发送停止读取, 停止读取时切记要发送停止信号
    50  	ReadCh(conn conn.Connection) (<-chan *readerRes, chan<- interface{})
    51  }
    52  
    53  type defaultReader struct{}
    54  
    55  func (d *defaultReader) ReadCh(conn conn.Connection) (<-chan *readerRes, chan<- interface{}) {
    56  	c := make(chan *readerRes, 5)
    57  	done := make(chan interface{})
    58  
    59  	go func() {
    60  		defer func() {
    61  			e := recover()
    62  			if e != nil {
    63  				logger.E("error on runRead msg from connection %v", e)
    64  			}
    65  		}()
    66  		for {
    67  			select {
    68  			case <-done:
    69  				goto CLOSE
    70  			default:
    71  				m, err := d.Read(conn)
    72  				res := recyclePool.Get().(*readerRes)
    73  				if err != nil {
    74  					res.err = err
    75  					c <- res
    76  					if messages.IsDecodeError(err) {
    77  						continue
    78  					}
    79  					goto CLOSE
    80  				} else {
    81  					res.m = m
    82  					c <- res
    83  				}
    84  			}
    85  		}
    86  	CLOSE:
    87  		close(c)
    88  	}()
    89  	return c, done
    90  }
    91  
    92  func (d *defaultReader) Read(conn conn.Connection) (*messages.GlideMessage, error) {
    93  	// TODO 2021-12-3 校验数据包
    94  	bytes, err := conn.Read()
    95  	if err != nil {
    96  		return nil, err
    97  	}
    98  	m := messages.NewEmptyMessage()
    99  	err = codec.Decode(bytes, m)
   100  	return m, err
   101  }