From d98b1b4438230571bd7fbc5aaa89843e92937bd3 Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@golang.org>
Date: Wed, 2 Oct 2019 20:37:15 +0000
Subject: [PATCH] http2: fix memory leak in random write scheduler

In certain shutdown cases (from the client and/or server), the http2
Server can Push stream-specific frames on closed streams. This caused
memory leaks in the random write scheduler.

As a conservative fix for backporting, just clear the map element
whenever its queue value is empty. The map entry is re-created as
needed anyway. This isn't perfectly ideal (it adds a map+delete and
free queue put+get) in the case where a stream is open & actively
writing, but it's an easy fix for now. A future CL can optimize all
this code. It looks like there are some other good optimization
opportunities in related code anyway. But I'd rather that happen on
master and not be done in a backported change.

Fixes golang/go#33812

Change-Id: I21508ba2ebc361e8b8532d0d1cebf882e82c473c
Reviewed-on: https://go-review.googlesource.com/c/net/+/198462
Reviewed-by: Bryan C. Mills <bcmills@google.com>
---
 http2/writesched_random.go      |  9 +++++++--
 http2/writesched_random_test.go | 16 ++++++++++++++++
 http2/writesched_test.go        |  5 +++++
 3 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/http2/writesched_random.go b/http2/writesched_random.go
index 36d7919f..9a7b9e58 100644
--- a/http2/writesched_random.go
+++ b/http2/writesched_random.go
@@ -19,7 +19,8 @@ type randomWriteScheduler struct {
 	zero writeQueue
 
 	// sq contains the stream-specific queues, keyed by stream ID.
-	// When a stream is idle or closed, it's deleted from the map.
+	// When a stream is idle, closed, or emptied, it's deleted
+	// from the map.
 	sq map[uint32]*writeQueue
 
 	// pool of empty queues for reuse.
@@ -63,8 +64,12 @@ func (ws *randomWriteScheduler) Pop() (FrameWriteRequest, bool) {
 		return ws.zero.shift(), true
 	}
 	// Iterate over all non-idle streams until finding one that can be consumed.
-	for _, q := range ws.sq {
+	for streamID, q := range ws.sq {
 		if wr, ok := q.consume(math.MaxInt32); ok {
+			if q.empty() {
+				delete(ws.sq, streamID)
+				ws.queuePool.put(q)
+			}
 			return wr, true
 		}
 	}
diff --git a/http2/writesched_random_test.go b/http2/writesched_random_test.go
index 3bf4aa36..1f501b4b 100644
--- a/http2/writesched_random_test.go
+++ b/http2/writesched_random_test.go
@@ -41,4 +41,20 @@ func TestRandomScheduler(t *testing.T) {
 			t.Errorf("frame not found for stream %d", id)
 		}
 	}
+
+	// Verify that we clean up maps for empty queues in all cases (golang.org/issue/33812)
+	const arbitraryStreamID = 123
+	ws.Push(makeHandlerPanicRST(arbitraryStreamID))
+	rws := ws.(*randomWriteScheduler)
+	if got, want := len(rws.sq), 1; got != want {
+		t.Fatalf("len of 123 stream = %v; want %v", got, want)
+	}
+	_, ok := ws.Pop()
+	if !ok {
+		t.Fatal("expected to be able to Pop")
+	}
+	if got, want := len(rws.sq), 0; got != want {
+		t.Fatalf("len of 123 stream = %v; want %v", got, want)
+	}
+
 }
diff --git a/http2/writesched_test.go b/http2/writesched_test.go
index 0807056b..99be5a77 100644
--- a/http2/writesched_test.go
+++ b/http2/writesched_test.go
@@ -20,6 +20,11 @@ func makeWriteHeadersRequest(streamID uint32) FrameWriteRequest {
 	return FrameWriteRequest{&writeResHeaders{streamID: streamID, httpResCode: 200}, st, nil}
 }
 
+func makeHandlerPanicRST(streamID uint32) FrameWriteRequest {
+	st := &stream{id: streamID}
+	return FrameWriteRequest{&handlerPanicRST{StreamID: streamID}, st, nil}
+}
+
 func checkConsume(wr FrameWriteRequest, nbytes int32, want []FrameWriteRequest) error {
 	consumed, rest, n := wr.Consume(nbytes)
 	var wantConsumed, wantRest FrameWriteRequest