Repository: quic-go/quic-go

Issue:
probe packet PING packing logic can be simplified
It seems unnecessary to add a PING to the retransmission queue, just to dequeue it from the packet packer when packing the probe packer:
https://github.com/quic-go/quic-go/blob/ea6d198c3d71e2b4b865b886342455fd297debef/connection.go#L2080-L2087

Instead, the packet packer could just generate a fresh PING frame if there's no actual data to be sent. We don't need to involve the retransmission queue at all.

Hints from issue comments:
@marten-seemann: Blocked on #4823.

Solution patch:
diff --git a/connection.go b/connection.go
index 8ca6ae5c..bf261c30 100644
--- a/connection.go
+++ b/connection.go
@@ -2319,29 +2319,25 @@ func (s *connection) sendProbePacket(sendMode ackhandler.SendMode, now time.Time
 	// Queue probe packets until we actually send out a packet,
 	// or until there are no more packets to queue.
 	var packet *coalescedPacket
-	for {
+	for packet == nil {
 		if wasQueued := s.sentPacketHandler.QueueProbePacket(encLevel); !wasQueued {
 			break
 		}
 		var err error
-		packet, err = s.packer.MaybePackPTOProbePacket(encLevel, s.maxPacketSize(), now, s.version)
+		packet, err = s.packer.PackPTOProbePacket(encLevel, s.maxPacketSize(), false, now, s.version)
 		if err != nil {
 			return err
 		}
-		if packet != nil {
-			break
-		}
 	}
 	if packet == nil {
-		s.retransmissionQueue.AddPing(encLevel)
 		var err error
-		packet, err = s.packer.MaybePackPTOProbePacket(encLevel, s.maxPacketSize(), now, s.version)
+		packet, err = s.packer.PackPTOProbePacket(encLevel, s.maxPacketSize(), true, now, s.version)
 		if err != nil {
 			return err
 		}
 	}
 	if packet == nil || (len(packet.longHdrPackets) == 0 && packet.shortHdrPacket == nil) {
-		return fmt.Errorf("connection BUG: couldn't pack %s probe packet", encLevel)
+		return fmt.Errorf("connection BUG: couldn't pack %s probe packet: %v", encLevel, packet)
 	}
 	return s.sendPackedCoalescedPacket(packet, s.sentPacketHandler.ECNMode(packet.IsOnlyShortHeaderPacket()), now)
 }
diff --git a/packet_packer.go b/packet_packer.go
index 45c190b5..5cd5f623 100644
--- a/packet_packer.go
+++ b/packet_packer.go
@@ -21,8 +21,8 @@ var errNothingToPack = errors.New("nothing to pack")
 type packer interface {
 	PackCoalescedPacket(onlyAck bool, maxPacketSize protocol.ByteCount, now time.Time, v protocol.Version) (*coalescedPacket, error)
 	PackAckOnlyPacket(maxPacketSize protocol.ByteCount, now time.Time, v protocol.Version) (shortHeaderPacket, *packetBuffer, error)
-	AppendPacket(buf *packetBuffer, maxPacketSize protocol.ByteCount, now time.Time, v protocol.Version) (shortHeaderPacket, error)
-	MaybePackPTOProbePacket(protocol.EncryptionLevel, protocol.ByteCount, time.Time, protocol.Version) (*coalescedPacket, error)
+	AppendPacket(_ *packetBuffer, maxPacketSize protocol.ByteCount, now time.Time, v protocol.Version) (shortHeaderPacket, error)
+	PackPTOProbePacket(_ protocol.EncryptionLevel, _ protocol.ByteCount, addPingIfEmpty bool, now time.Time, v protocol.Version) (*coalescedPacket, error)
 	PackConnectionClose(*qerr.TransportError, protocol.ByteCount, protocol.Version) (*coalescedPacket, error)
 	PackApplicationClose(*qerr.ApplicationError, protocol.ByteCount, protocol.Version) (*coalescedPacket, error)
 	PackPathProbePacket(protocol.ConnectionID, []ackhandler.Frame, protocol.Version) (shortHeaderPacket, *packetBuffer, error)
@@ -348,6 +348,7 @@ func (p *packetPacker) PackCoalescedPacket(onlyAck bool, maxSize protocol.ByteCo
 			maxSize-protocol.ByteCount(initialSealer.Overhead()),
 			protocol.EncryptionInitial,
 			now,
+			false,
 			onlyAck,
 			true,
 			v,
@@ -370,6 +371,7 @@ func (p *packetPacker) PackCoalescedPacket(onlyAck bool, maxSize protocol.ByteCo
 				maxSize-size-protocol.ByteCount(handshakeSealer.Overhead()),
 				protocol.EncryptionHandshake,
 				now,
+				false,
 				onlyAck,
 				size == 0,
 				v,
@@ -497,6 +499,7 @@ func (p *packetPacker) maybeGetCryptoPacket(
 	maxPacketSize protocol.ByteCount,
 	encLevel protocol.EncryptionLevel,
 	now time.Time,
+	addPingIfEmpty bool,
 	onlyAck, ackAllowed bool,
 	v protocol.Version,
 ) (*wire.ExtendedHeader, payload) {
@@ -530,12 +533,17 @@ func (p *packetPacker) maybeGetCryptoPacket(
 	if ackAllowed {
 		ack = p.acks.GetAckFrame(encLevel, now, !hasRetransmission && !hasData)
 	}
+	var pl payload
 	if !hasData && !hasRetransmission && ack == nil {
-		// nothing to send
-		return nil, payload{}
+		if !addPingIfEmpty {
+			// nothing to send
+			return nil, payload{}
+		}
+		ping := &wire.PingFrame{}
+		pl.frames = append(pl.frames, ackhandler.Frame{Frame: ping, Handler: emptyHandler{}})
+		pl.length += ping.Length(v)
 	}
 
-	var pl payload
 	if ack != nil {
 		pl.ack = ack
 		pl.length = ack.Length(v)
@@ -565,7 +573,7 @@ func (p *packetPacker) maybeGetCryptoPacket(
 		}
 	} else if s.HasData() {
 		cf := s.PopCryptoFrame(maxPacketSize)
-		pl.frames = []ackhandler.Frame{{Frame: cf, Handler: handler}}
+		pl.frames = append(pl.frames, ackhandler.Frame{Frame: cf, Handler: handler})
 		pl.length += cf.Length(v)
 	}
 	return hdr, pl
@@ -709,33 +717,15 @@ func (p *packetPacker) composeNextPacket(
 	return pl
 }
 
-func (p *packetPacker) MaybePackPTOProbePacket(
+func (p *packetPacker) PackPTOProbePacket(
 	encLevel protocol.EncryptionLevel,
 	maxPacketSize protocol.ByteCount,
+	addPingIfEmpty bool,
 	now time.Time,
 	v protocol.Version,
 ) (*coalescedPacket, error) {
 	if encLevel == protocol.Encryption1RTT {
-		s, err := p.cryptoSetup.Get1RTTSealer()
-		if err != nil {
-			return nil, err
-		}
-		kp := s.KeyPhase()
-		connID := p.getDestConnID()
-		pn, pnLen := p.pnManager.PeekPacketNumber(protocol.Encryption1RTT)
-		hdrLen := wire.ShortHeaderLen(connID, pnLen)
-		pl := p.maybeGetAppDataPacket(maxPacketSize-protocol.ByteCount(s.Overhead())-hdrLen, false, true, now, v)
-		if pl.length == 0 {
-			return nil, nil
-		}
-		buffer := getPacketBuffer()
-		packet := &coalescedPacket{buffer: buffer}
-		shp, err := p.appendShortHeaderPacket(buffer, connID, pn, pnLen, kp, pl, 0, maxPacketSize, s, false, v)
-		if err != nil {
-			return nil, err
-		}
-		packet.shortHdrPacket = &shp
-		return packet, nil
+		return p.packPTOProbePacket1RTT(maxPacketSize, addPingIfEmpty, now, v)
 	}
 
 	var sealer handshake.LongHeaderSealer
@@ -756,7 +746,15 @@ func (p *packetPacker) MaybePackPTOProbePacket(
 	default:
 		panic("unknown encryption level")
 	}
-	hdr, pl := p.maybeGetCryptoPacket(maxPacketSize-protocol.ByteCount(sealer.Overhead()), encLevel, now, false, true, v)
+	hdr, pl := p.maybeGetCryptoPacket(
+		maxPacketSize-protocol.ByteCount(sealer.Overhead()),
+		encLevel,
+		now,
+		addPingIfEmpty,
+		false,
+		true,
+		v,
+	)
 	if pl.length == 0 {
 		return nil, nil
 	}
@@ -776,6 +774,34 @@ func (p *packetPacker) MaybePackPTOProbePacket(
 	return packet, nil
 }
 
+func (p *packetPacker) packPTOProbePacket1RTT(maxPacketSize protocol.ByteCount, addPingIfEmpty bool, now time.Time, v protocol.Version) (*coalescedPacket, error) {
+	s, err := p.cryptoSetup.Get1RTTSealer()
+	if err != nil {
+		return nil, err
+	}
+	kp := s.KeyPhase()
+	connID := p.getDestConnID()
+	pn, pnLen := p.pnManager.PeekPacketNumber(protocol.Encryption1RTT)
+	hdrLen := wire.ShortHeaderLen(connID, pnLen)
+	pl := p.maybeGetAppDataPacket(maxPacketSize-protocol.ByteCount(s.Overhead())-hdrLen, false, true, now, v)
+	if pl.length == 0 {
+		if !addPingIfEmpty {
+			return nil, nil
+		}
+		ping := &wire.PingFrame{}
+		pl.frames = append(pl.frames, ackhandler.Frame{Frame: ping, Handler: emptyHandler{}})
+		pl.length += ping.Length(v)
+	}
+	buffer := getPacketBuffer()
+	packet := &coalescedPacket{buffer: buffer}
+	shp, err := p.appendShortHeaderPacket(buffer, connID, pn, pnLen, kp, pl, 0, maxPacketSize, s, false, v)
+	if err != nil {
+		return nil, err
+	}
+	packet.shortHdrPacket = &shp
+	return packet, nil
+}
+
 func (p *packetPacker) PackMTUProbePacket(ping ackhandler.Frame, size protocol.ByteCount, v protocol.Version) (shortHeaderPacket, *packetBuffer, error) {
 	pl := payload{
 		frames: []ackhandler.Frame{ping},
@@ -983,3 +1009,10 @@ func (p *packetPacker) encryptPacket(raw []byte, sealer sealer, pn protocol.Pack
 func (p *packetPacker) SetToken(token []byte) {
 	p.token = token
 }
+
+type emptyHandler struct{}
+
+var _ ackhandler.FrameHandler = emptyHandler{}
+
+func (emptyHandler) OnAcked(wire.Frame) {}
+func (emptyHandler) OnLost(wire.Frame)  {}
diff --git a/retransmission_queue.go b/retransmission_queue.go
index 14ae1e3b..8c78edbf 100644
--- a/retransmission_queue.go
+++ b/retransmission_queue.go
@@ -23,22 +23,6 @@ func newRetransmissionQueue() *retransmissionQueue {
 	return &retransmissionQueue{}
 }
 
-// AddPing queues a ping.
-// It is used when a probe packet needs to be sent
-func (q *retransmissionQueue) AddPing(encLevel protocol.EncryptionLevel) {
-	//nolint:exhaustive // Cannot send probe packets for 0-RTT.
-	switch encLevel {
-	case protocol.EncryptionInitial:
-		q.addInitial(&wire.PingFrame{})
-	case protocol.EncryptionHandshake:
-		q.addHandshake(&wire.PingFrame{})
-	case protocol.Encryption1RTT:
-		q.addAppData(&wire.PingFrame{})
-	default:
-		panic("unexpected encryption level")
-	}
-}
-
 func (q *retransmissionQueue) addInitial(f wire.Frame) {
 	if cf, ok := f.(*wire.CryptoFrame); ok {
 		q.initialCryptoData = append(q.initialCryptoData, cf)



Test patch:
diff --git a/connection_test.go b/connection_test.go
index ee4547d1..caaa8ff5 100644
--- a/connection_test.go
+++ b/connection_test.go
@@ -2295,8 +2295,8 @@ func testConnectionPTOProbePackets(t *testing.T, encLevel protocol.EncryptionLev
 	sph.EXPECT().QueueProbePacket(encLevel).Return(false)
 	sph.EXPECT().SentPacket(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any())
 
-	tc.packer.EXPECT().MaybePackPTOProbePacket(encLevel, gomock.Any(), gomock.Any(), protocol.Version1).DoAndReturn(
-		func(encLevel protocol.EncryptionLevel, maxSize protocol.ByteCount, t time.Time, version protocol.Version) (*coalescedPacket, error) {
+	tc.packer.EXPECT().PackPTOProbePacket(encLevel, gomock.Any(), true, gomock.Any(), protocol.Version1).DoAndReturn(
+		func(protocol.EncryptionLevel, protocol.ByteCount, bool, time.Time, protocol.Version) (*coalescedPacket, error) {
 			return &coalescedPacket{
 				buffer:         getPacketBuffer(),
 				shortHdrPacket: &shortHeaderPacket{PacketNumber: 1},
diff --git a/mock_packer_test.go b/mock_packer_test.go
index a1c6ed3c..106f11d4 100644
--- a/mock_packer_test.go
+++ b/mock_packer_test.go
@@ -44,18 +44,18 @@ func (m *MockPacker) EXPECT() *MockPackerMockRecorder {
 }
 
 // AppendPacket mocks base method.
-func (m *MockPacker) AppendPacket(buf *packetBuffer, maxPacketSize protocol.ByteCount, now time.Time, v protocol.Version) (shortHeaderPacket, error) {
+func (m *MockPacker) AppendPacket(arg0 *packetBuffer, maxPacketSize protocol.ByteCount, now time.Time, v protocol.Version) (shortHeaderPacket, error) {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "AppendPacket", buf, maxPacketSize, now, v)
+	ret := m.ctrl.Call(m, "AppendPacket", arg0, maxPacketSize, now, v)
 	ret0, _ := ret[0].(shortHeaderPacket)
 	ret1, _ := ret[1].(error)
 	return ret0, ret1
 }
 
 // AppendPacket indicates an expected call of AppendPacket.
-func (mr *MockPackerMockRecorder) AppendPacket(buf, maxPacketSize, now, v any) *MockPackerAppendPacketCall {
+func (mr *MockPackerMockRecorder) AppendPacket(arg0, maxPacketSize, now, v any) *MockPackerAppendPacketCall {
 	mr.mock.ctrl.T.Helper()
-	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "AppendPacket", reflect.TypeOf((*MockPacker)(nil).AppendPacket), buf, maxPacketSize, now, v)
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "AppendPacket", reflect.TypeOf((*MockPacker)(nil).AppendPacket), arg0, maxPacketSize, now, v)
 	return &MockPackerAppendPacketCall{Call: call}
 }
 
@@ -82,45 +82,6 @@ func (c *MockPackerAppendPacketCall) DoAndReturn(f func(*packetBuffer, protocol.
 	return c
 }
 
-// MaybePackPTOProbePacket mocks base method.
-func (m *MockPacker) MaybePackPTOProbePacket(arg0 protocol.EncryptionLevel, arg1 protocol.ByteCount, arg2 time.Time, arg3 protocol.Version) (*coalescedPacket, error) {
-	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "MaybePackPTOProbePacket", arg0, arg1, arg2, arg3)
-	ret0, _ := ret[0].(*coalescedPacket)
-	ret1, _ := ret[1].(error)
-	return ret0, ret1
-}
-
-// MaybePackPTOProbePacket indicates an expected call of MaybePackPTOProbePacket.
-func (mr *MockPackerMockRecorder) MaybePackPTOProbePacket(arg0, arg1, arg2, arg3 any) *MockPackerMaybePackPTOProbePacketCall {
-	mr.mock.ctrl.T.Helper()
-	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MaybePackPTOProbePacket", reflect.TypeOf((*MockPacker)(nil).MaybePackPTOProbePacket), arg0, arg1, arg2, arg3)
-	return &MockPackerMaybePackPTOProbePacketCall{Call: call}
-}
-
-// MockPackerMaybePackPTOProbePacketCall wrap *gomock.Call
-type MockPackerMaybePackPTOProbePacketCall struct {
-	*gomock.Call
-}
-
-// Return rewrite *gomock.Call.Return
-func (c *MockPackerMaybePackPTOProbePacketCall) Return(arg0 *coalescedPacket, arg1 error) *MockPackerMaybePackPTOProbePacketCall {
-	c.Call = c.Call.Return(arg0, arg1)
-	return c
-}
-
-// Do rewrite *gomock.Call.Do
-func (c *MockPackerMaybePackPTOProbePacketCall) Do(f func(protocol.EncryptionLevel, protocol.ByteCount, time.Time, protocol.Version) (*coalescedPacket, error)) *MockPackerMaybePackPTOProbePacketCall {
-	c.Call = c.Call.Do(f)
-	return c
-}
-
-// DoAndReturn rewrite *gomock.Call.DoAndReturn
-func (c *MockPackerMaybePackPTOProbePacketCall) DoAndReturn(f func(protocol.EncryptionLevel, protocol.ByteCount, time.Time, protocol.Version) (*coalescedPacket, error)) *MockPackerMaybePackPTOProbePacketCall {
-	c.Call = c.Call.DoAndReturn(f)
-	return c
-}
-
 // PackAckOnlyPacket mocks base method.
 func (m *MockPacker) PackAckOnlyPacket(maxPacketSize protocol.ByteCount, now time.Time, v protocol.Version) (shortHeaderPacket, *packetBuffer, error) {
 	m.ctrl.T.Helper()
@@ -318,6 +279,45 @@ func (c *MockPackerPackMTUProbePacketCall) DoAndReturn(f func(ackhandler.Frame,
 	return c
 }
 
+// PackPTOProbePacket mocks base method.
+func (m *MockPacker) PackPTOProbePacket(arg0 protocol.EncryptionLevel, arg1 protocol.ByteCount, addPingIfEmpty bool, now time.Time, v protocol.Version) (*coalescedPacket, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "PackPTOProbePacket", arg0, arg1, addPingIfEmpty, now, v)
+	ret0, _ := ret[0].(*coalescedPacket)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// PackPTOProbePacket indicates an expected call of PackPTOProbePacket.
+func (mr *MockPackerMockRecorder) PackPTOProbePacket(arg0, arg1, addPingIfEmpty, now, v any) *MockPackerPackPTOProbePacketCall {
+	mr.mock.ctrl.T.Helper()
+	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "PackPTOProbePacket", reflect.TypeOf((*MockPacker)(nil).PackPTOProbePacket), arg0, arg1, addPingIfEmpty, now, v)
+	return &MockPackerPackPTOProbePacketCall{Call: call}
+}
+
+// MockPackerPackPTOProbePacketCall wrap *gomock.Call
+type MockPackerPackPTOProbePacketCall struct {
+	*gomock.Call
+}
+
+// Return rewrite *gomock.Call.Return
+func (c *MockPackerPackPTOProbePacketCall) Return(arg0 *coalescedPacket, arg1 error) *MockPackerPackPTOProbePacketCall {
+	c.Call = c.Call.Return(arg0, arg1)
+	return c
+}
+
+// Do rewrite *gomock.Call.Do
+func (c *MockPackerPackPTOProbePacketCall) Do(f func(protocol.EncryptionLevel, protocol.ByteCount, bool, time.Time, protocol.Version) (*coalescedPacket, error)) *MockPackerPackPTOProbePacketCall {
+	c.Call = c.Call.Do(f)
+	return c
+}
+
+// DoAndReturn rewrite *gomock.Call.DoAndReturn
+func (c *MockPackerPackPTOProbePacketCall) DoAndReturn(f func(protocol.EncryptionLevel, protocol.ByteCount, bool, time.Time, protocol.Version) (*coalescedPacket, error)) *MockPackerPackPTOProbePacketCall {
+	c.Call = c.Call.DoAndReturn(f)
+	return c
+}
+
 // PackPathProbePacket mocks base method.
 func (m *MockPacker) PackPathProbePacket(arg0 protocol.ConnectionID, arg1 []ackhandler.Frame, arg2 protocol.Version) (shortHeaderPacket, *packetBuffer, error) {
 	m.ctrl.T.Helper()
diff --git a/packet_packer_test.go b/packet_packer_test.go
index ed412ced..c588408f 100644
--- a/packet_packer_test.go
+++ b/packet_packer_test.go
@@ -809,7 +809,7 @@ func testPackProbePacket(t *testing.T, encLevel protocol.EncryptionLevel, perspe
 	tp.pnManager.EXPECT().PeekPacketNumber(encLevel).Return(protocol.PacketNumber(0x42), protocol.PacketNumberLen2)
 	tp.pnManager.EXPECT().PopPacketNumber(encLevel).Return(protocol.PacketNumber(0x42))
 
-	p, err := tp.packer.MaybePackPTOProbePacket(encLevel, maxPacketSize, time.Now(), protocol.Version1)
+	p, err := tp.packer.PackPTOProbePacket(encLevel, maxPacketSize, false, time.Now(), protocol.Version1)
 	require.NoError(t, err)
 	require.NotNil(t, p)
 	require.Len(t, p.longHdrPackets, 1)
@@ -832,17 +832,6 @@ func testPackProbePacket(t *testing.T, encLevel protocol.EncryptionLevel, perspe
 	require.Empty(t, more)
 }
 
-func TestPackProbePacketNothingToSend(t *testing.T) {
-	mockCtrl := gomock.NewController(t)
-	tp := newTestPacketPacker(t, mockCtrl, protocol.PerspectiveClient)
-	tp.sealingManager.EXPECT().GetInitialSealer().Return(newMockShortHeaderSealer(mockCtrl), nil)
-	tp.ackFramer.EXPECT().GetAckFrame(protocol.EncryptionInitial, gomock.Any(), true)
-
-	p, err := tp.packer.MaybePackPTOProbePacket(protocol.EncryptionInitial, protocol.MaxByteCount, time.Now(), protocol.Version1)
-	require.NoError(t, err)
-	require.Nil(t, p)
-}
-
 func TestPack1RTTProbePacket(t *testing.T) {
 	const maxPacketSize protocol.ByteCount = 999
 
@@ -861,7 +850,7 @@ func TestPack1RTTProbePacket(t *testing.T) {
 		},
 	)
 
-	p, err := tp.packer.MaybePackPTOProbePacket(protocol.Encryption1RTT, maxPacketSize, time.Now(), protocol.Version1)
+	p, err := tp.packer.PackPTOProbePacket(protocol.Encryption1RTT, maxPacketSize, false, time.Now(), protocol.Version1)
 	require.NoError(t, err)
 	require.NotNil(t, p)
 	require.True(t, p.IsOnlyShortHeaderPacket())
@@ -873,18 +862,61 @@ func TestPack1RTTProbePacket(t *testing.T) {
 	require.Equal(t, maxPacketSize, packet.Length)
 }
 
-func TestPackProbePacketNothingToPack(t *testing.T) {
+func TestPackPTOProbePacketNothingToPack(t *testing.T) {
+	t.Run("Initial", func(t *testing.T) {
+		testPackPTOProbePacketNothingToPack(t, protocol.EncryptionInitial)
+	})
+	t.Run("Handshake", func(t *testing.T) {
+		testPackPTOProbePacketNothingToPack(t, protocol.EncryptionHandshake)
+	})
+	t.Run("1-RTT", func(t *testing.T) {
+		testPackPTOProbePacketNothingToPack(t, protocol.Encryption1RTT)
+	})
+}
+
+func testPackPTOProbePacketNothingToPack(t *testing.T, encLevel protocol.EncryptionLevel) {
+	const maxPacketSize protocol.ByteCount = 1234
 	mockCtrl := gomock.NewController(t)
 	tp := newTestPacketPacker(t, mockCtrl, protocol.PerspectiveServer)
 
-	tp.sealingManager.EXPECT().Get1RTTSealer().Return(newMockShortHeaderSealer(mockCtrl), nil)
-	tp.pnManager.EXPECT().PeekPacketNumber(protocol.Encryption1RTT).Return(protocol.PacketNumber(0x42), protocol.PacketNumberLen2)
-	tp.ackFramer.EXPECT().GetAckFrame(protocol.Encryption1RTT, gomock.Any(), true)
-	tp.framer.EXPECT().HasData()
+	switch encLevel {
+	case protocol.EncryptionInitial:
+		tp.sealingManager.EXPECT().GetInitialSealer().Return(newMockShortHeaderSealer(mockCtrl), nil).Times(2)
+	case protocol.EncryptionHandshake:
+		tp.sealingManager.EXPECT().GetHandshakeSealer().Return(newMockShortHeaderSealer(mockCtrl), nil).Times(2)
+	case protocol.Encryption1RTT:
+		tp.sealingManager.EXPECT().Get1RTTSealer().Return(newMockShortHeaderSealer(mockCtrl), nil).Times(2)
+		tp.framer.EXPECT().HasData().Times(2)
+	}
+	tp.pnManager.EXPECT().PeekPacketNumber(encLevel).Return(protocol.PacketNumber(0x42), protocol.PacketNumberLen2).MaxTimes(2)
+	tp.ackFramer.EXPECT().GetAckFrame(encLevel, gomock.Any(), true).Times(2)
 
-	packet, err := tp.packer.MaybePackPTOProbePacket(protocol.Encryption1RTT, protocol.MaxByteCount, time.Now(), protocol.Version1)
+	// don't force a PING to be sent
+	packet, err := tp.packer.PackPTOProbePacket(encLevel, maxPacketSize, false, time.Now(), protocol.Version1)
 	require.NoError(t, err)
 	require.Nil(t, packet)
+
+	// now force a PING to be sent
+	tp.pnManager.EXPECT().PopPacketNumber(encLevel).Return(protocol.PacketNumber(0x42))
+	packet, err = tp.packer.PackPTOProbePacket(encLevel, maxPacketSize, true, time.Now(), protocol.Version1)
+	require.NoError(t, err)
+	require.NotNil(t, packet)
+	var frames []ackhandler.Frame
+	switch encLevel {
+	case protocol.EncryptionInitial, protocol.EncryptionHandshake:
+		require.Len(t, packet.longHdrPackets, 1)
+		require.Nil(t, packet.shortHdrPacket)
+		require.Equal(t, encLevel, packet.longHdrPackets[0].EncryptionLevel())
+		frames = packet.longHdrPackets[0].frames
+	case protocol.Encryption1RTT:
+		require.Empty(t, packet.longHdrPackets)
+		require.NotNil(t, packet.shortHdrPacket)
+		frames = packet.shortHdrPacket.Frames
+	}
+
+	require.Len(t, frames, 1)
+	require.Equal(t, &wire.PingFrame{}, frames[0].Frame)
+	require.Equal(t, emptyHandler{}, frames[0].Handler)
 }
 
 func TestPackMTUProbePacket(t *testing.T) {
diff --git a/retransmission_queue_test.go b/retransmission_queue_test.go
index c0132848..3a3ae83b 100644
--- a/retransmission_queue_test.go
+++ b/retransmission_queue_test.go
@@ -96,12 +96,6 @@ var _ = Describe("Retransmission queue", func() {
 			Expect(q.HasInitialData()).To(BeTrue())
 			Expect(q.GetInitialFrame(protocol.MaxByteCount, protocol.Version1)).To(Equal(f))
 		})
-
-		It("adds a PING", func() {
-			q.AddPing(protocol.EncryptionInitial)
-			Expect(q.HasInitialData()).To(BeTrue())
-			Expect(q.GetInitialFrame(protocol.MaxByteCount, protocol.Version1)).To(Equal(&wire.PingFrame{}))
-		})
 	})
 
 	Context("Handshake data", func() {
@@ -185,12 +179,6 @@ var _ = Describe("Retransmission queue", func() {
 			Expect(q.HasHandshakeData()).To(BeTrue())
 			Expect(q.GetHandshakeFrame(protocol.MaxByteCount, protocol.Version1)).To(Equal(f))
 		})
-
-		It("adds a PING", func() {
-			q.AddPing(protocol.EncryptionHandshake)
-			Expect(q.HasHandshakeData()).To(BeTrue())
-			Expect(q.GetHandshakeFrame(protocol.MaxByteCount, protocol.Version1)).To(Equal(&wire.PingFrame{}))
-		})
 	})
 
 	Context("Application data", func() {
@@ -214,11 +202,5 @@ var _ = Describe("Retransmission queue", func() {
 			Expect(q.HasAppData()).To(BeTrue())
 			Expect(q.GetAppDataFrame(protocol.MaxByteCount, protocol.Version1)).To(Equal(f))
 		})
-
-		It("adds a PING", func() {
-			q.AddPing(protocol.Encryption1RTT)
-			Expect(q.HasAppData()).To(BeTrue())
-			Expect(q.GetAppDataFrame(protocol.MaxByteCount, protocol.Version1)).To(Equal(&wire.PingFrame{}))
-		})
 	})
 })