From bd6d893cc4ae10ac5e82f206d4420eb70b23e597 Mon Sep 17 00:00:00 2001
From: Lucas Clemente <lucas@clemente.io>
Date: Thu, 19 May 2016 14:16:42 +0200
Subject: [PATCH] replace time.After in session with a single timer

fixes #128, fixes #127
---
 session.go      | 49 +++++++++++++++++++++++++++++++++----------------
 session_test.go | 13 -------------
 2 files changed, 33 insertions(+), 29 deletions(-)

diff --git a/session.go b/session.go
index 06810f53..9b4f14d3 100644
--- a/session.go
+++ b/session.go
@@ -73,6 +73,10 @@ type Session struct {
 	// Used to calculate the next packet number from the truncated wire
 	// representation, and sent back in public reset packets
 	lastRcvdPacketNumber protocol.PacketNumber
+
+	lastNetworkActivityTime time.Time
+
+	timer *time.Timer
 }
 
 // newSession makes a new session
@@ -97,6 +101,8 @@ func newSession(conn connection, v protocol.VersionNumber, connectionID protocol
 		connectionParametersManager: connectionParametersManager,
 		undecryptablePackets:        make([]receivedPacket, 0, protocol.MaxUndecryptablePackets),
 		aeadChanged:                 make(chan struct{}, 1),
+		timer:                       time.NewTimer(0),
+		lastNetworkActivityTime: time.Now(),
 	}
 
 	cryptoStream, _ := session.OpenStream(1)
@@ -128,39 +134,39 @@ func (s *Session) run() {
 		default:
 		}
 
-		// Note: receive at a nil channel blocks forever
+		// Calculate the minimum of all timeouts
 
-		var smallPacketSendTimer <-chan time.Time
+		now := time.Now()
+		firstTimeout := utils.InfDuration
+		// Small packet send delay
 		if !s.smallPacketDelayedOccurranceTime.IsZero() {
-			smallPacketSendTimer = time.After(time.Now().Sub(s.smallPacketDelayedOccurranceTime))
+			firstTimeout = utils.MinDuration(firstTimeout, s.smallPacketDelayedOccurranceTime.Add(protocol.SmallPacketSendDelay).Sub(now))
 		}
+		// RTOs
+		firstTimeout = utils.MinDuration(firstTimeout, s.sentPacketHandler.TimeToFirstRTO())
+		// Idle connection timeout
+		firstTimeout = utils.MinDuration(firstTimeout, s.lastNetworkActivityTime.Add(s.connectionParametersManager.GetIdleConnectionStateLifetime()).Sub(now))
 
-		var rtoTimer <-chan time.Time
-		if d := s.sentPacketHandler.TimeToFirstRTO(); d != utils.InfDuration {
-			rtoTimer = time.After(d)
-		}
+		s.timer.Reset(firstTimeout)
 
 		var err error
 		select {
 		case <-s.closeChan:
 			return
+		case <-s.timer.C:
+			// We do all the interesting stuff after the switch statement, so
+			// nothing to see here.
+		case <-s.sendingScheduled:
+			// We do all the interesting stuff after the switch statement, so
+			// nothing to see here.
 		case p := <-s.receivedPackets:
 			err = s.handlePacketImpl(p.remoteAddr, p.publicHeader, p.data)
 			if qErr, ok := err.(*qerr.QuicError); ok && qErr.ErrorCode == qerr.DecryptionFailure {
 				s.tryQueueingUndecryptablePacket(p)
 				continue
 			}
-			s.scheduleSending()
-		case <-s.sendingScheduled:
-			err = s.maybeSendPacket()
-		case <-smallPacketSendTimer:
-			err = s.sendPacket()
-		case <-rtoTimer:
-			err = s.sendPacket()
 		case <-s.aeadChanged:
 			s.tryDecryptingQueuedPackets()
-		case <-time.After(s.connectionParametersManager.GetIdleConnectionStateLifetime()):
-			s.Close(qerr.Error(qerr.NetworkIdleTimeout, "No recent network activity."))
 		}
 
 		if err != nil {
@@ -177,11 +183,18 @@ func (s *Session) run() {
 			}
 		}
 
+		if err := s.maybeSendPacket(); err != nil {
+			s.Close(err)
+		}
+		if time.Now().Sub(s.lastNetworkActivityTime) > s.connectionParametersManager.GetIdleConnectionStateLifetime() {
+			s.Close(qerr.Error(qerr.NetworkIdleTimeout, "No recent network activity."))
+		}
 		s.garbageCollectStreams()
 	}
 }
 
 func (s *Session) handlePacketImpl(remoteAddr interface{}, hdr *publicHeader, data []byte) error {
+	s.lastNetworkActivityTime = time.Now()
 	r := bytes.NewReader(data)
 
 	// Calculate packet number
@@ -419,6 +432,10 @@ func (s *Session) maybeSendPacket() error {
 		s.smallPacketDelayedOccurranceTime = time.Now()
 	}
 
+	if time.Now().Sub(s.smallPacketDelayedOccurranceTime) > protocol.SmallPacketSendDelay {
+		return s.sendPacket()
+	}
+
 	return nil
 }
 
diff --git a/session_test.go b/session_test.go
index aff121c2..7623bb2f 100644
--- a/session_test.go
+++ b/session_test.go
@@ -414,19 +414,6 @@ var _ = Describe("Session", func() {
 			Expect(session.sendingScheduled).To(Receive())
 		})
 
-		It("sends after receiving a packet", func() {
-			Expect(session.sendingScheduled).NotTo(Receive())
-			session.receivedPackets <- receivedPacket{
-				publicHeader: &publicHeader{},
-				data: []byte{
-					// FNV hash + "foobar"
-					0x18, 0x6f, 0x44, 0xba, 0x97, 0x35, 0xd, 0x6f, 0xbf, 0x64, 0x3c, 0x79, 0x66, 0x6f, 0x6f, 0x62, 0x61, 0x72,
-				},
-			}
-			session.run()
-			Expect(session.sendingScheduled).To(Receive())
-		})
-
 		Context("bundling of small packets", func() {
 			It("bundles two small frames into one packet", func() {
 				go session.run()