From 4cb8bf3101f39bde0bbe04fc774a03e6749937aa Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Wed, 4 Sep 2019 20:36:45 +0700
Subject: [PATCH] put STREAM frames back into the pool when they are
 acknowledged

---
 internal/ackhandler/frame.go                    | 2 +-
 internal/ackhandler/sent_packet_handler.go      | 2 +-
 internal/ackhandler/sent_packet_handler_test.go | 7 ++++++-
 send_stream.go                                  | 4 +++-
 send_stream_test.go                             | 8 ++++----
 5 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/internal/ackhandler/frame.go b/internal/ackhandler/frame.go
index 5731c2bc9..aed6038d9 100644
--- a/internal/ackhandler/frame.go
+++ b/internal/ackhandler/frame.go
@@ -5,5 +5,5 @@ import "github.com/lucas-clemente/quic-go/internal/wire"
 type Frame struct {
 	wire.Frame // nil if the frame has already been acknowledged in another packet
 	OnLost     func(wire.Frame)
-	OnAcked    func()
+	OnAcked    func(wire.Frame)
 }
diff --git a/internal/ackhandler/sent_packet_handler.go b/internal/ackhandler/sent_packet_handler.go
index 4b4815637..f542db324 100644
--- a/internal/ackhandler/sent_packet_handler.go
+++ b/internal/ackhandler/sent_packet_handler.go
@@ -446,7 +446,7 @@ func (h *sentPacketHandler) onPacketAcked(p *Packet, rcvTime time.Time) error {
 
 	for _, f := range p.Frames {
 		if f.OnAcked != nil {
-			f.OnAcked()
+			f.OnAcked(f.Frame)
 		}
 	}
 	if p.includedInBytesInFlight {
diff --git a/internal/ackhandler/sent_packet_handler_test.go b/internal/ackhandler/sent_packet_handler_test.go
index af1ee4a8c..0859a0a76 100644
--- a/internal/ackhandler/sent_packet_handler_test.go
+++ b/internal/ackhandler/sent_packet_handler_test.go
@@ -202,9 +202,14 @@ var _ = Describe("SentPacketHandler", func() {
 
 			It("calls the OnAcked callback", func() {
 				var acked bool
+				ping := &wire.PingFrame{}
 				handler.SentPacket(ackElicitingPacket(&Packet{
 					PacketNumber: 13,
-					Frames:       []Frame{{Frame: &wire.PingFrame{}, OnAcked: func() { acked = true }}},
+					Frames: []Frame{{Frame: ping, OnAcked: func(f wire.Frame) {
+						Expect(f).To(Equal(ping))
+						acked = true
+					},
+					}},
 				}))
 				ack := &wire.AckFrame{AckRanges: []wire.AckRange{{Smallest: 13, Largest: 13}}}
 				Expect(handler.ReceivedAck(ack, 1, protocol.Encryption1RTT, time.Now())).To(Succeed())
diff --git a/send_stream.go b/send_stream.go
index a2a3ad2a2..c2e1aa7d0 100644
--- a/send_stream.go
+++ b/send_stream.go
@@ -270,7 +270,9 @@ func (s *sendStream) getDataForWriting(f *wire.StreamFrame, maxBytes protocol.By
 	f.FinBit = s.finishedWriting && s.dataForWriting == nil && !s.finSent
 }
 
-func (s *sendStream) frameAcked() {
+func (s *sendStream) frameAcked(f wire.Frame) {
+	f.(*wire.StreamFrame).PutBack()
+
 	s.mutex.Lock()
 	s.numOutstandingFrames--
 	if s.numOutstandingFrames < 0 {
diff --git a/send_stream_test.go b/send_stream_test.go
index 3f47b0def..e15eb5f36 100644
--- a/send_stream_test.go
+++ b/send_stream_test.go
@@ -786,7 +786,7 @@ var _ = Describe("Send Stream", func() {
 			// Acknowledge all frames.
 			// We don't expect the stream to be completed, since we still need to send the FIN.
 			for _, f := range frames {
-				f.OnAcked()
+				f.OnAcked(f.Frame)
 			}
 
 			// Now close the stream and acknowledge the FIN.
@@ -794,7 +794,7 @@ var _ = Describe("Send Stream", func() {
 			frame, _ := str.popStreamFrame(protocol.MaxByteCount)
 			Expect(frame).ToNot(BeNil())
 			mockSender.EXPECT().onStreamCompleted(streamID)
-			frame.OnAcked()
+			frame.OnAcked(frame.Frame)
 		})
 
 		It("doesn't say it's completed when there are frames waiting to be retransmitted", func() {
@@ -824,7 +824,7 @@ var _ = Describe("Send Stream", func() {
 
 			// lose the first frame, acknowledge all others
 			for _, f := range frames[1:] {
-				f.OnAcked()
+				f.OnAcked(f.Frame)
 			}
 			frames[0].OnLost(frames[0].Frame)
 
@@ -832,7 +832,7 @@ var _ = Describe("Send Stream", func() {
 			ret, _ := str.popStreamFrame(protocol.MaxByteCount)
 			Expect(ret).ToNot(BeNil())
 			mockSender.EXPECT().onStreamCompleted(streamID)
-			ret.OnAcked()
+			ret.OnAcked(ret.Frame)
 		})
 	})
 })