From 6a80d50fb93f10d7f2292386ac80c5915c8019bf Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Sun, 11 Mar 2018 12:13:29 +0700
Subject: [PATCH] optimize updating the RTT when receiving an ACK

---
 internal/ackhandler/sent_packet_handler.go    | 19 +++++--------------
 internal/ackhandler/sent_packet_history.go    |  7 +++++++
 .../ackhandler/sent_packet_history_test.go    | 10 ++++++++++
 3 files changed, 22 insertions(+), 14 deletions(-)

diff --git a/internal/ackhandler/sent_packet_handler.go b/internal/ackhandler/sent_packet_handler.go
index cc6bbd7d1..df13b8f2a 100644
--- a/internal/ackhandler/sent_packet_handler.go
+++ b/internal/ackhandler/sent_packet_handler.go
@@ -239,20 +239,11 @@ func (h *sentPacketHandler) determineNewlyAckedPackets(ackFrame *wire.AckFrame)
 }
 
 func (h *sentPacketHandler) maybeUpdateRTT(largestAcked protocol.PacketNumber, ackDelay time.Duration, rcvTime time.Time) bool {
-	var rttUpdated bool
-	h.packetHistory.Iterate(func(p *Packet) (bool, error) {
-		if p.PacketNumber == largestAcked {
-			h.rttStats.UpdateRTT(rcvTime.Sub(p.sendTime), ackDelay, rcvTime)
-			rttUpdated = true
-			return false, nil
-		}
-		// Packets are sorted by number, so we can stop searching
-		if p.PacketNumber > largestAcked {
-			return false, nil
-		}
-		return true, nil
-	})
-	return rttUpdated
+	if p := h.packetHistory.GetPacket(largestAcked); p != nil {
+		h.rttStats.UpdateRTT(rcvTime.Sub(p.sendTime), ackDelay, rcvTime)
+		return true
+	}
+	return false
 }
 
 func (h *sentPacketHandler) updateLossDetectionAlarm(now time.Time) {
diff --git a/internal/ackhandler/sent_packet_history.go b/internal/ackhandler/sent_packet_history.go
index eeb247094..486daf7a2 100644
--- a/internal/ackhandler/sent_packet_history.go
+++ b/internal/ackhandler/sent_packet_history.go
@@ -23,6 +23,13 @@ func (h *sentPacketHistory) SentPacket(p *Packet) {
 	h.packetMap[p.PacketNumber] = el
 }
 
+func (h *sentPacketHistory) GetPacket(p protocol.PacketNumber) *Packet {
+	if el, ok := h.packetMap[p]; ok {
+		return &el.Value
+	}
+	return nil
+}
+
 // Iterate iterates through all packets.
 // The callback must not modify the history.
 func (h *sentPacketHistory) Iterate(cb func(*Packet) (cont bool, err error)) error {
diff --git a/internal/ackhandler/sent_packet_history_test.go b/internal/ackhandler/sent_packet_history_test.go
index a88101fa7..7026f3ea3 100644
--- a/internal/ackhandler/sent_packet_history_test.go
+++ b/internal/ackhandler/sent_packet_history_test.go
@@ -53,6 +53,16 @@ var _ = Describe("SentPacketHistory", func() {
 		Expect(front.PacketNumber).To(Equal(protocol.PacketNumber(2)))
 	})
 
+	It("gets a packet by packet number", func() {
+		p := &Packet{PacketNumber: 2}
+		hist.SentPacket(p)
+		Expect(hist.GetPacket(2)).To(Equal(p))
+	})
+
+	It("returns nil if the packet doesn't exist", func() {
+		Expect(hist.GetPacket(1337)).To(BeNil())
+	})
+
 	It("removes packets", func() {
 		hist.SentPacket(&Packet{PacketNumber: 1})
 		hist.SentPacket(&Packet{PacketNumber: 4})