From 0e69b373e2203cf7740b6a5980f9b54caaaa53cf Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Sat, 27 Aug 2022 18:29:57 +0300
Subject: [PATCH] use a sync.Pool to avoid allocations of ackhandler.Packet

---
 internal/ackhandler/interfaces.go          | 20 -----------
 internal/ackhandler/packet.go              | 42 ++++++++++++++++++++++
 internal/ackhandler/sent_packet_handler.go |  3 ++
 packet_packer.go                           | 19 +++++-----
 4 files changed, 55 insertions(+), 29 deletions(-)
 create mode 100644 internal/ackhandler/packet.go

diff --git a/internal/ackhandler/interfaces.go b/internal/ackhandler/interfaces.go
index 226bfcbbcc5..9079fb0be24 100644
--- a/internal/ackhandler/interfaces.go
+++ b/internal/ackhandler/interfaces.go
@@ -7,26 +7,6 @@ import (
 	"github.com/lucas-clemente/quic-go/internal/wire"
 )
 
-// A Packet is a packet
-type Packet struct {
-	PacketNumber    protocol.PacketNumber
-	Frames          []Frame
-	LargestAcked    protocol.PacketNumber // InvalidPacketNumber if the packet doesn't contain an ACK
-	Length          protocol.ByteCount
-	EncryptionLevel protocol.EncryptionLevel
-	SendTime        time.Time
-
-	IsPathMTUProbePacket bool // We don't report the loss of Path MTU probe packets to the congestion controller.
-
-	includedInBytesInFlight bool
-	declaredLost            bool
-	skippedPacket           bool
-}
-
-func (p *Packet) outstanding() bool {
-	return !p.declaredLost && !p.skippedPacket && !p.IsPathMTUProbePacket
-}
-
 // SentPacketHandler handles ACKs received for outgoing packets
 type SentPacketHandler interface {
 	// SentPacket may modify the packet
diff --git a/internal/ackhandler/packet.go b/internal/ackhandler/packet.go
new file mode 100644
index 00000000000..0095b18bf0a
--- /dev/null
+++ b/internal/ackhandler/packet.go
@@ -0,0 +1,42 @@
+package ackhandler
+
+import (
+	"sync"
+	"time"
+
+	"github.com/lucas-clemente/quic-go/internal/protocol"
+)
+
+// A Packet is a packet
+type Packet struct {
+	PacketNumber    protocol.PacketNumber
+	Frames          []Frame
+	LargestAcked    protocol.PacketNumber // InvalidPacketNumber if the packet doesn't contain an ACK
+	Length          protocol.ByteCount
+	EncryptionLevel protocol.EncryptionLevel
+	SendTime        time.Time
+
+	IsPathMTUProbePacket bool // We don't report the loss of Path MTU probe packets to the congestion controller.
+
+	includedInBytesInFlight bool
+	declaredLost            bool
+	skippedPacket           bool
+}
+
+var packetPool = sync.Pool{New: func() any { return &Packet{} }}
+
+func GetPacket() *Packet {
+	p := packetPool.Get().(*Packet)
+	p.includedInBytesInFlight = false
+	p.declaredLost = false
+	p.skippedPacket = false
+	return p
+}
+
+// We currently only return Packets back into the pool when they're acknowledged (not when they're lost).
+// This simplifies the code, and gives the vast majority of the performance benefit we can gain from using the pool.
+func putPacket(p *Packet) { packetPool.Put(p) }
+
+func (p *Packet) outstanding() bool {
+	return !p.declaredLost && !p.skippedPacket && !p.IsPathMTUProbePacket
+}
diff --git a/internal/ackhandler/sent_packet_handler.go b/internal/ackhandler/sent_packet_handler.go
index 5a8cd70e8fd..9a7f023a6b7 100644
--- a/internal/ackhandler/sent_packet_handler.go
+++ b/internal/ackhandler/sent_packet_handler.go
@@ -334,7 +334,10 @@ func (h *sentPacketHandler) ReceivedAck(ack *wire.AckFrame, encLevel protocol.En
 			acked1RTTPacket = true
 		}
 		h.removeFromBytesInFlight(p)
+		putPacket(p)
 	}
+	// After this point, we must not use ackedPackets any longer!
+	// We've already returned the buffers.
 
 	// Reset the pto_count unless the client is unsure if the server has validated the client's address.
 	if h.peerCompletedAddressValidation {
diff --git a/packet_packer.go b/packet_packer.go
index aec485f0050..a860122893d 100644
--- a/packet_packer.go
+++ b/packet_packer.go
@@ -100,15 +100,16 @@ func (p *packetContents) ToAckHandlerPacket(now time.Time, q *retransmissionQueu
 			p.frames[i].OnLost = q.AddAppData
 		}
 	}
-	return &ackhandler.Packet{
-		PacketNumber:         p.header.PacketNumber,
-		LargestAcked:         largestAcked,
-		Frames:               p.frames,
-		Length:               p.length,
-		EncryptionLevel:      encLevel,
-		SendTime:             now,
-		IsPathMTUProbePacket: p.isMTUProbePacket,
-	}
+
+	ap := ackhandler.GetPacket()
+	ap.PacketNumber = p.header.PacketNumber
+	ap.LargestAcked = largestAcked
+	ap.Frames = p.frames
+	ap.Length = p.length
+	ap.EncryptionLevel = encLevel
+	ap.SendTime = now
+	ap.IsPathMTUProbePacket = p.isMTUProbePacket
+	return ap
 }
 
 func getMaxPacketSize(addr net.Addr) protocol.ByteCount {