From 105964bfac15e9e0100aaa22a5c433bcd30c4c3c Mon Sep 17 00:00:00 2001
From: Jakob Buchgraber <buchgr@google.com>
Date: Mon, 6 Apr 2015 17:53:34 -0700
Subject: [PATCH] Fix memory leak by adding the Http2StreamRemovalPolicy to the
 channel pipeline.

In benchmarks we would see grpc talking up tens of gigabytes of memory. A heap dump
revealed that streams would not get cleaned up and stick around in memory forever.
---
 .../main/java/io/grpc/transport/netty/Http2Negotiator.java  | 6 ++++--
 .../java/io/grpc/transport/netty/NettyClientTransport.java  | 4 ++--
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/netty/src/main/java/io/grpc/transport/netty/Http2Negotiator.java b/netty/src/main/java/io/grpc/transport/netty/Http2Negotiator.java
index a20b61dd81..9d20d9867c 100644
--- a/netty/src/main/java/io/grpc/transport/netty/Http2Negotiator.java
+++ b/netty/src/main/java/io/grpc/transport/netty/Http2Negotiator.java
@@ -160,7 +160,8 @@ public class Http2Negotiator {
   /**
    * Create a plaintext upgrade negotiation for HTTP/1.1 to HTTP/2.
    */
-  public static Negotiation plaintextUpgrade(final Http2ConnectionHandler handler) {
+  public static Negotiation plaintextUpgrade(final ChannelHandler streamRemovalPolicy,
+                                             Http2ConnectionHandler handler) {
     // Register the plaintext upgrader
     Http2ClientUpgradeCodec upgradeCodec = new Http2ClientUpgradeCodec(handler);
     HttpClientCodec httpClientCodec = new HttpClientCodec();
@@ -170,6 +171,7 @@ public class Http2Negotiator {
     final ChannelInitializer<Channel> initializer = new ChannelInitializer<Channel>() {
       @Override
       public void initChannel(Channel ch) throws Exception {
+        ch.pipeline().addLast(streamRemovalPolicy);
         ch.pipeline().addLast(upgrader);
         ch.pipeline().addLast(completionHandler);
       }
@@ -200,7 +202,7 @@ public class Http2Negotiator {
   /**
    * Create a "no-op" negotiation that simply assumes the protocol to already be negotiated.
    */
-  public static Negotiation plaintext(final ChannelHandler handler) {
+  public static Negotiation plaintext(final ChannelHandler... handlers) {
     final ChannelInitializer<Channel> initializer = new ChannelInitializer<Channel>() {
       @Override
       public void initChannel(Channel ch) throws Exception {
diff --git a/netty/src/main/java/io/grpc/transport/netty/NettyClientTransport.java b/netty/src/main/java/io/grpc/transport/netty/NettyClientTransport.java
index 5af9a0abf4..65acf64802 100644
--- a/netty/src/main/java/io/grpc/transport/netty/NettyClientTransport.java
+++ b/netty/src/main/java/io/grpc/transport/netty/NettyClientTransport.java
@@ -137,11 +137,11 @@ class NettyClientTransport implements ClientTransport {
     handler = newHandler(streamRemovalPolicy);
     switch (negotiationType) {
       case PLAINTEXT:
-        negotiation = Http2Negotiator.plaintext(handler);
+        negotiation = Http2Negotiator.plaintext(streamRemovalPolicy, handler);
         ssl = false;
         break;
       case PLAINTEXT_UPGRADE:
-        negotiation = Http2Negotiator.plaintextUpgrade(handler);
+        negotiation = Http2Negotiator.plaintextUpgrade(streamRemovalPolicy, handler);
         ssl = false;
         break;
       case TLS: