counter

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