diff --git a/cart-packet-queue.go b/cart-packet-queue.go index 11585ba..bca122f 100644 --- a/cart-packet-queue.go +++ b/cart-packet-queue.go @@ -14,7 +14,7 @@ type CartPacketQueue struct { const CurrentPacketVersion = 2 -type CartListener map[CartId]chan []byte +type CartListener map[CartId]Listener func NewCartPacketQueue(connection net.Conn) *CartPacketQueue { @@ -58,13 +58,16 @@ func (p *CartPacketQueue) HandleConnection(connection net.Conn) error { func (p *CartPacketQueue) HandleData(t uint32, id CartId, data []byte) { p.mu.Lock() defer p.mu.Unlock() - l, ok := p.expectedPackages[t] + pl, ok := p.expectedPackages[t] if ok { - ch, ok := (*l)[id] + l, ok := (*pl)[id] if ok { - ch <- data - close(ch) - delete(*l, id) + l.Chan <- data + l.Count-- + if l.Count == 0 { + close(l.Chan) + delete(*pl, id) + } } } data = nil @@ -75,17 +78,24 @@ func (p *CartPacketQueue) Expect(messageType uint32, id CartId) <-chan []byte { defer p.mu.Unlock() l, ok := p.expectedPackages[messageType] if ok { - if ch, idOk := (*l)[id]; idOk { - return ch + if idl, idOk := (*l)[id]; idOk { + idl.Count++ + return idl.Chan } ch := make(chan []byte) - (*l)[id] = ch + (*l)[id] = Listener{ + Chan: ch, + Count: 1, + } return ch } ch := make(chan []byte) p.expectedPackages[messageType] = &CartListener{ - id: ch, + id: Listener{ + Chan: ch, + Count: 1, + }, } return ch