123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103 |
- package msg
- import (
- "io"
- "reflect"
- )
- func AsyncHandler(f func(Message)) func(Message) {
- return func(m Message) {
- go f(m)
- }
- }
- type Dispatcher struct {
- rw io.ReadWriter
- sendCh chan Message
- doneCh chan struct{}
- msgHandlers map[reflect.Type]func(Message)
- defaultHandler func(Message)
- }
- func NewDispatcher(rw io.ReadWriter) *Dispatcher {
- return &Dispatcher{
- rw: rw,
- sendCh: make(chan Message, 100),
- doneCh: make(chan struct{}),
- msgHandlers: make(map[reflect.Type]func(Message)),
- }
- }
- func (d *Dispatcher) Run() {
- go d.sendLoop()
- go d.readLoop()
- }
- func (d *Dispatcher) sendLoop() {
- for {
- select {
- case <-d.doneCh:
- return
- case m := <-d.sendCh:
- _ = WriteMsg(d.rw, m)
- }
- }
- }
- func (d *Dispatcher) readLoop() {
- for {
- m, err := ReadMsg(d.rw)
- if err != nil {
- close(d.doneCh)
- return
- }
- if handler, ok := d.msgHandlers[reflect.TypeOf(m)]; ok {
- handler(m)
- } else if d.defaultHandler != nil {
- d.defaultHandler(m)
- }
- }
- }
- func (d *Dispatcher) Send(m Message) error {
- select {
- case <-d.doneCh:
- return io.EOF
- case d.sendCh <- m:
- return nil
- }
- }
- func (d *Dispatcher) SendChannel() chan Message {
- return d.sendCh
- }
- func (d *Dispatcher) RegisterHandler(msg Message, handler func(Message)) {
- d.msgHandlers[reflect.TypeOf(msg)] = handler
- }
- func (d *Dispatcher) RegisterDefaultHandler(handler func(Message)) {
- d.defaultHandler = handler
- }
- func (d *Dispatcher) Done() chan struct{} {
- return d.doneCh
- }
|