put STREAM frames back into the pool when they are acknowledged

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)
 }

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 {

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())

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 {

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)
 		})
 	})
 })