From 659f98125586022459e59420f2372db50cb5028c Mon Sep 17 00:00:00 2001
From: Mahak Mukhi <mmukhi@google.com>
Date: Thu, 9 Mar 2017 20:03:59 -0800
Subject: [PATCH] more initial commits

---
 transport/http2_server.go   | 29 +++++++++++++++++++++++++----
 transport/transport_test.go | 10 ++++++++--
 2 files changed, 33 insertions(+), 6 deletions(-)

diff --git a/transport/http2_server.go b/transport/http2_server.go
index deb91d2b..1a6735eb 100644
--- a/transport/http2_server.go
+++ b/transport/http2_server.go
@@ -768,16 +768,38 @@ func (t *http2Server) keepalive() {
 	maxIdle := time.NewTimer(t.kp.MaxConnectionIdle)
 	maxAge := time.NewTimer(t.kp.MaxConnectionAge)
 	keepalive := time.NewTimer(t.kp.Time)
+	// NOTE: All exit paths of this function should reset their
+	// respecitve timers. A failure to do so will cause the
+	// following clean-up to deadlock and eventually leak.
+	defer func() {
+		if !maxIdle.Stop() {
+			<-maxIdle.C
+		}
+		if !maxAge.Stop() {
+			<-maxAge.C
+		}
+		if !keepalive.Stop() {
+			<-keepalive.C
+		}
+	}()
 	t.mu.Lock()
-	idle := t.idle
+	oidle := t.idle
 	t.mu.Unlock()
 	for {
 		select {
 		case <-maxIdle.C:
-			if idle == t.idle {
-				// send go away
+			t.mu.Lock()
+			idle := t.idle
+			t.mu.Unlock()
+			if idle == oidle {
+				t.Drain()
+				//TODO: Graceful connection termination!?
+				// Reseting the timer so that the clean-up doesn't deadlock.
+				maxIdle.Reset(infinity)
+				// New Stream might get created by the time GoAway is written on the wire.
 				continue
 			}
+			oidle = idle
 			if idle.IsZero() {
 				maxIdle.Reset(t.kp.MaxConnectionIdle)
 				continue
@@ -786,7 +808,6 @@ func (t *http2Server) keepalive() {
 		case <-maxAge.C:
 		case <-keepalive.C:
 		case <-t.shutdownChan:
-			// TODO(mmukhi): clean-up
 			return
 		}
 	}
diff --git a/transport/transport_test.go b/transport/transport_test.go
index fcbd566c..0d3c6226 100644
--- a/transport/transport_test.go
+++ b/transport/transport_test.go
@@ -303,7 +303,7 @@ func setUpWithNoPingServer(t *testing.T, copts ConnectOptions, done chan net.Con
 // MaxConnectionIdle time.
 func TestMaxConnectionIdle(t *testing.T) {
 	serverConfig := &ServerConfig{
-		keepaliveParams: keepalive.ServerParams{
+		keepaliveParams: keepalive.ServerParameters{
 			MaxConnectionIdle: 2 * time.Second,
 		},
 	}
@@ -317,8 +317,14 @@ func TestMaxConnectionIdle(t *testing.T) {
 	stream.mu.Lock()
 	stream.rstStream = true
 	stream.mu.Unlock()
-	client.CloseStream()
+	client.CloseStream(stream, nil)
 	// wait for server to see that closed stream and max age to send goaway after no new RPCs are mode
+	timeout := time.NewTimer(time.Second * 4)
+	select {
+	case <-client.GoAway():
+	case <-timeout.C:
+		t.Fatalf("Test timed out, expected a GoAway from server")
+	}
 }
 
 func TestKeepaliveClientClosesIdleTransport(t *testing.T) {