make sure that all packets in the send queue are sent before closing

send_queue.go

@@ -1,16 +1,18 @@
 package quic
 
 type sendQueue struct {
-	queue     chan *packedPacket
-	closeChan chan struct{}
-	conn      connection
+	queue       chan *packedPacket
+	closeCalled chan struct{} // runStopped when Close() is called
+	runStopped  chan struct{} // runStopped when the run loop returns
+	conn        connection
 }
 
 func newSendQueue(conn connection) *sendQueue {
 	s := &sendQueue{
-		conn:      conn,
-		closeChan: make(chan struct{}),
-		queue:     make(chan *packedPacket, 1),
+		conn:        conn,
+		runStopped:  make(chan struct{}),
+		closeCalled: make(chan struct{}),
+		queue:       make(chan *packedPacket, 1),
 	}
 	return s
 }
@@ -20,20 +22,28 @@ func (h *sendQueue) Send(p *packedPacket) {
 }
 
 func (h *sendQueue) Run() error {
-	var p *packedPacket
+	defer close(h.runStopped)
+	var shouldClose bool
 	for {
-		select {
-		case <-h.closeChan:
+		if shouldClose && len(h.queue) == 0 {
 			return nil
-		case p = <-h.queue:
 		}
-		if err := h.conn.Write(p.raw); err != nil {
-			return err
+		select {
+		case <-h.closeCalled:
+			h.closeCalled = nil // prevent this case from being selected again
+			// make sure that all queued packets are actually sent out
+			shouldClose = true
+		case p := <-h.queue:
+			if err := h.conn.Write(p.raw); err != nil {
+				return err
+			}
+			p.buffer.Release()
 		}
-		p.buffer.Release()
 	}
 }
 
 func (h *sendQueue) Close() {
-	close(h.closeChan)
+	close(h.closeCalled)
+	// wait until the run loop returned
+	<-h.runStopped
 }