update queue for max performance

packet-queue.go

@@ -1,98 +1,97 @@
 package main
 
 import (
-	"fmt"
 	"io"
 	"log"
 	"net"
 	"sync"
-	"time"
 )
 
-type PacketWithData struct {
-	MessageType uint32
-	Added       time.Time
-	Consumed    bool
-	Data        []byte
-}
+// type PacketWithData struct {
+// 	MessageType uint32
+// 	Added       time.Time
+// 	Consumed    bool
+// 	Data        []byte
+// }
 
 type PacketQueue struct {
-	mu      sync.RWMutex
-	Packets []PacketWithData
+	mu               sync.RWMutex
+	expectedPackages map[uint32]*Listener
+	//Packets []PacketWithData
 	//connection net.Conn
 }
 
-const cap = 150
+//const cap = 150
+
+type Listener struct {
+	Count int
+	Chan  chan []byte
+}
 
 func NewPacketQueue(connection net.Conn) *PacketQueue {
 
 	queue := &PacketQueue{
-		Packets: make([]PacketWithData, 0, cap),
+		expectedPackages: make(map[uint32]*Listener),
+		//Packets: make([]PacketWithData, 0, cap+1),
 		//connection: connection,
 	}
-	go func() {
-		defer connection.Close()
-		var packet Packet
-		for {
-			err := ReadPacket(connection, &packet)
-			if err != nil {
-				if err == io.EOF {
-					return
-				}
-				log.Printf("Error receiving packet: %v\n", err)
-				//return
-			}
-			data, err := GetPacketData(connection, packet.DataLength)
-			if err != nil {
-				log.Printf("Error receiving packet data: %v\n", err)
-			}
-
-			go queue.HandleData(packet.MessageType, data)
-		}
-	}()
+	go queue.HandleConnection(connection)
 	return queue
 }
 
-func (p *PacketQueue) HandleData(t uint32, data []byte) {
-	ts := time.Now()
-	l := make([]PacketWithData, 0, cap)
-	p.mu.RLock()
-	breakAt := ts.Add(-time.Millisecond * 250)
-	for _, packet := range p.Packets {
-		if !packet.Consumed && packet.Added.After(breakAt) {
-			l = append(l, packet)
-			if len(l) >= cap {
-				break
-			}
-		}
-	}
-	p.mu.RUnlock()
-	p.mu.Lock()
-	p.Packets = append([]PacketWithData{
-		{
-			MessageType: t,
-			Added:       ts,
-			Data:        data,
-		},
-	}, l...)
-	p.mu.Unlock()
-}
-
-func (p *PacketQueue) Expect(messageType uint32, timeToWait time.Duration) (*PacketWithData, error) {
-	start := time.Now().Add(-time.Millisecond)
+func (p *PacketQueue) HandleConnection(connection net.Conn) error {
 
+	defer connection.Close()
+	var packet Packet
 	for {
-		if time.Since(start) > timeToWait {
-			return nil, fmt.Errorf("timeout waiting for message type %d", messageType)
-		}
-		p.mu.RLock()
-		defer p.mu.RUnlock()
-		for _, packet := range p.Packets {
-			if !packet.Consumed && packet.MessageType == messageType && packet.Added.After(start) {
-				packet.Consumed = true
-				return &packet, nil
+		err := ReadPacket(connection, &packet)
+		if err != nil {
+			if err == io.EOF {
+				return nil
 			}
+			log.Printf("Error receiving packet: %v\n", err)
+			return err
 		}
-		time.Sleep(time.Millisecond * 4)
+		data, err := GetPacketData(connection, packet.DataLength)
+		if err != nil {
+			log.Printf("Error receiving packet data: %v\n", err)
+			return err
+		}
+		go p.HandleData(packet.MessageType, data)
 	}
 }
+
+func (p *PacketQueue) HandleData(t uint32, data []byte) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	l, ok := p.expectedPackages[t]
+	if ok {
+		l.Chan <- data
+		l.Count--
+		if l.Count == 0 {
+			close(l.Chan)
+			delete(p.expectedPackages, t)
+		}
+		return
+	}
+	data = nil
+}
+
+func (p *PacketQueue) Expect(messageType uint32) <-chan []byte {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	l, ok := p.expectedPackages[messageType]
+	if ok {
+		l.Count++
+		return l.Chan
+	}
+
+	ch := make(chan []byte)
+	p.expectedPackages[messageType] = &Listener{
+		Count: 1,
+		Chan:  ch,
+	}
+
+	return ch
+
+}