immediately remove reset tokens when retiring a connection ID

conn_id_manager.go

@@ -30,15 +30,13 @@ type connIDManager struct {
 
 	addStatelessResetToken    func(protocol.StatelessResetToken)
 	removeStatelessResetToken func(protocol.StatelessResetToken)
-	retireStatelessResetToken func(protocol.StatelessResetToken)
 	queueControlFrame         func(wire.Frame)
 }
 
 func newConnIDManager(
 	initialDestConnID protocol.ConnectionID,
 	addStatelessResetToken func(protocol.StatelessResetToken),
 	removeStatelessResetToken func(protocol.StatelessResetToken),
-	retireStatelessResetToken func(protocol.StatelessResetToken),
 	queueControlFrame func(wire.Frame),
 ) *connIDManager {
 	b := make([]byte, 8)
@@ -48,7 +46,6 @@ func newConnIDManager(
 		activeConnectionID:        initialDestConnID,
 		addStatelessResetToken:    addStatelessResetToken,
 		removeStatelessResetToken: removeStatelessResetToken,
-		retireStatelessResetToken: retireStatelessResetToken,
 		queueControlFrame:         queueControlFrame,
 		rand:                      mrand.New(mrand.NewSource(seed)),
 	}
@@ -150,7 +147,7 @@ func (h *connIDManager) updateConnectionID() {
 	})
 	h.highestRetired = utils.MaxUint64(h.highestRetired, h.activeSequenceNumber)
 	if h.activeStatelessResetToken != nil {
-		h.retireStatelessResetToken(*h.activeStatelessResetToken)
+		h.removeStatelessResetToken(*h.activeStatelessResetToken)
 	}
 
 	front := h.queue.Remove(h.queue.Front())

conn_id_manager_test.go

@@ -12,21 +12,18 @@ var _ = Describe("Connection ID Manager", func() {
 		m             *connIDManager
 		frameQueue    []wire.Frame
 		tokenAdded    *protocol.StatelessResetToken
-		retiredTokens []protocol.StatelessResetToken
 		removedTokens []protocol.StatelessResetToken
 	)
 	initialConnID := protocol.ConnectionID{0, 0, 0, 0}
 
 	BeforeEach(func() {
 		frameQueue = nil
 		tokenAdded = nil
-		retiredTokens = nil
 		removedTokens = nil
 		m = newConnIDManager(
 			initialConnID,
 			func(token protocol.StatelessResetToken) { tokenAdded = &token },
 			func(token protocol.StatelessResetToken) { removedTokens = append(removedTokens, token) },
-			func(token protocol.StatelessResetToken) { retiredTokens = append(retiredTokens, token) },
 			func(f wire.Frame,
 			) {
 				frameQueue = append(frameQueue, f)
@@ -282,8 +279,8 @@ var _ = Describe("Connection ID Manager", func() {
 			if !connID.Equal(lastConnID) {
 				counter++
 				lastConnID = connID
-				Expect(retiredTokens).To(HaveLen(1))
-				retiredTokens = nil
+				Expect(removedTokens).To(HaveLen(1))
+				removedTokens = nil
 				Expect(m.Add(&wire.NewConnectionIDFrame{
 					SequenceNumber:      uint64(s),
 					ConnectionID:        protocol.ConnectionID{s, s, s, s},
@@ -344,13 +341,12 @@ var _ = Describe("Connection ID Manager", func() {
 			ConnectionID:   protocol.ConnectionID{1, 3, 3, 7},
 		})).To(Succeed())
 		Expect(m.Get()).To(Equal(protocol.ConnectionID{2, 2, 2, 2}))
-		Expect(retiredTokens).To(HaveLen(1))
-		Expect(retiredTokens[0]).To(Equal(protocol.StatelessResetToken{1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1}))
+		Expect(removedTokens).To(HaveLen(1))
+		Expect(removedTokens[0]).To(Equal(protocol.StatelessResetToken{1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1}))
 	})
 
 	It("removes the currently active stateless reset token when it is closed", func() {
 		m.Close()
-		Expect(retiredTokens).To(BeEmpty())
 		Expect(removedTokens).To(BeEmpty())
 		Expect(m.Add(&wire.NewConnectionIDFrame{
 			SequenceNumber:      1,
@@ -360,7 +356,6 @@ var _ = Describe("Connection ID Manager", func() {
 		m.SetHandshakeComplete()
 		Expect(m.Get()).To(Equal(protocol.ConnectionID{1, 2, 3, 4}))
 		m.Close()
-		Expect(retiredTokens).To(BeEmpty())
 		Expect(removedTokens).To(HaveLen(1))
 		Expect(removedTokens[0]).To(Equal(protocol.StatelessResetToken{16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1}))
 	})

packet_handler_map.go

@@ -222,14 +222,6 @@ func (h *packetHandlerMap) RemoveResetToken(token protocol.StatelessResetToken)
 	h.mutex.Unlock()
 }
 
-func (h *packetHandlerMap) RetireResetToken(token protocol.StatelessResetToken) {
-	time.AfterFunc(h.deleteRetiredSessionsAfter, func() {
-		h.mutex.Lock()
-		delete(h.resetTokens, token)
-		h.mutex.Unlock()
-	})
-}
-
 func (h *packetHandlerMap) SetServer(s unknownPacketHandler) {
 	h.mutex.Lock()
 	h.server = s

packet_handler_map_test.go

@@ -322,20 +322,19 @@ var _ = Describe("Packet Handler Map", func() {
 					Eventually(destroyed).Should(BeClosed())
 				})
 
-				It("retires reset tokens", func() {
-					handler.deleteRetiredSessionsAfter = scaleDuration(10 * time.Millisecond)
+				It("removes reset tokens", func() {
 					connID := protocol.ConnectionID{0xde, 0xad, 0xbe, 0xef, 0x42}
 					packetHandler := NewMockPacketHandler(mockCtrl)
 					handler.Add(connID, packetHandler)
 					token := protocol.StatelessResetToken{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}
 					handler.AddResetToken(token, NewMockPacketHandler(mockCtrl))
-					handler.RetireResetToken(token)
+					handler.RemoveResetToken(token)
+					// don't EXPECT any call to packetHandler.destroy()
 					packetHandler.EXPECT().handlePacket(gomock.Any())
 					p := append([]byte{0x40} /* short header packet */, connID.Bytes()...)
 					p = append(p, make([]byte, 50)...)
 					p = append(p, token[:]...)
 
-					time.Sleep(scaleDuration(30 * time.Millisecond))
 					handler.handlePacket(&receivedPacket{data: p})
 				})
 

session.go

@@ -88,7 +88,6 @@ type sessionRunner interface {
 	ReplaceWithClosed(protocol.ConnectionID, packetHandler)
 	AddResetToken(protocol.StatelessResetToken, packetHandler)
 	RemoveResetToken(protocol.StatelessResetToken)
-	RetireResetToken(protocol.StatelessResetToken)
 }
 
 type handshakeRunner struct {
@@ -258,7 +257,6 @@ var newSession = func(
 		destConnID,
 		func(token protocol.StatelessResetToken) { runner.AddResetToken(token, s) },
 		runner.RemoveResetToken,
-		runner.RetireResetToken,
 		s.queueControlFrame,
 	)
 	s.connIDGenerator = newConnIDGenerator(
@@ -382,7 +380,6 @@ var newClientSession = func(
 		destConnID,
 		func(token protocol.StatelessResetToken) { runner.AddResetToken(token, s) },
 		runner.RemoveResetToken,
-		runner.RetireResetToken,
 		s.queueControlFrame,
 	)
 	s.connIDGenerator = newConnIDGenerator(