Fix Inprocess memory leak

census/src/test/java/io/grpc/census/CensusModulesTest.java

@@ -56,6 +56,7 @@
 import io.grpc.ClientInterceptors;
 import io.grpc.ClientStreamTracer;
 import io.grpc.Context;
+import io.grpc.KnownLength;
 import io.grpc.Metadata;
 import io.grpc.MethodDescriptor;
 import io.grpc.ServerCall;
@@ -99,6 +100,7 @@
 import io.opencensus.trace.Tracer;
 import io.opencensus.trace.propagation.BinaryFormat;
 import io.opencensus.trace.propagation.SpanContextParseException;
+import java.io.IOException;
 import java.io.InputStream;
 import java.util.HashSet;
 import java.util.List;
@@ -136,7 +138,7 @@ public class CensusModulesTest {
       ClientStreamTracer.StreamInfo.newBuilder()
           .setCallOptions(CallOptions.DEFAULT.withOption(NAME_RESOLUTION_DELAYED, 10L)).build();
 
-  private static class StringInputStream extends InputStream {
+  private static class StringInputStream extends InputStream implements KnownLength {
     final String string;
 
     StringInputStream(String string) {
@@ -149,6 +151,11 @@ public int read() {
       // passed to the InProcess server and consumed by MARSHALLER.parse().
       throw new UnsupportedOperationException("Should not be called");
     }
+
+    @Override
+    public int available() throws IOException {
+      return string == null ? 0 : string.length();
+    }
   }
 
   private static final MethodDescriptor.Marshaller<String> MARSHALLER =

core/src/testFixtures/java/io/grpc/internal/AbstractTransportTest.java

@@ -59,6 +59,7 @@
 import io.grpc.Status;
 import io.grpc.internal.testing.TestClientStreamTracer;
 import io.grpc.internal.testing.TestServerStreamTracer;
+import io.grpc.internal.testing.TestStreamTracer;
 import java.io.ByteArrayInputStream;
 import java.io.IOException;
 import java.io.InputStream;
@@ -249,6 +250,11 @@ protected long fakeCurrentTimeNanos() {
     throw new UnsupportedOperationException();
   }
 
+  protected void assertInProcessTransportAssumedMessageSize(
+          TestStreamTracer streamTracerSender, TestStreamTracer streamTracerReceiver) {
+    // implemented by SizesReportedInProcessTransportTest
+  }
+
   // TODO(ejona):
   //   multiple streams on same transport
   //   multiple client transports to same server
@@ -867,6 +873,7 @@ public void basicStream() throws Exception {
     assertThat(serverStreamTracer1.nextInboundEvent())
         .matches("inboundMessageRead\\(0, -?[0-9]+, -?[0-9]+\\)");
 
+    assertInProcessTransportAssumedMessageSize(clientStreamTracer1, serverStreamTracer1);
     Metadata serverHeaders = new Metadata();
     serverHeaders.put(asciiKey, "server");
     serverHeaders.put(asciiKey, "dupvalue");
@@ -903,6 +910,7 @@ public void basicStream() throws Exception {
         .matches("inboundMessageRead\\(0, -?[0-9]+, -?[0-9]+\\)");
     assertThat(clientStreamTracer1.getInboundWireSize()).isGreaterThan(0L);
     assertThat(clientStreamTracer1.getInboundUncompressedSize()).isGreaterThan(0L);
+    assertInProcessTransportAssumedMessageSize(serverStreamTracer1, clientStreamTracer1);
 
     message.close();
     assertNull("no additional message expected", clientStreamListener.messageQueue.poll());
@@ -1266,6 +1274,7 @@ public void onReady() {
     assertThat(clientStreamTracer1.getInboundUncompressedSize()).isGreaterThan(0L);
     assertThat(serverStreamTracer1.getOutboundWireSize()).isGreaterThan(0L);
     assertThat(serverStreamTracer1.getOutboundUncompressedSize()).isGreaterThan(0L);
+    assertInProcessTransportAssumedMessageSize(serverStreamTracer1, clientStreamTracer1);
     assertNull(clientStreamTracer1.getInboundTrailers());
     assertSame(status, clientStreamTracer1.getStatus());
     // There is a race between client cancelling and server closing.  The final status seen by the

inprocess/src/main/java/io/grpc/inprocess/InProcessChannelBuilder.java

@@ -239,7 +239,7 @@ static final class InProcessClientTransportFactory implements ClientTransportFac
     private final int maxInboundMetadataSize;
     private boolean closed;
     private final boolean includeCauseWithStatus;
-    long assumedMessageSize = -1;
+    private final long assumedMessageSize = -1;
 
     private InProcessClientTransportFactory(
         @Nullable ScheduledExecutorService scheduledExecutorService,
@@ -288,9 +288,5 @@ public void close() {
     public Collection<Class<? extends SocketAddress>> getSupportedSocketAddressTypes() {
       return Arrays.asList(InProcessSocketAddress.class, AnonymousInProcessSocketAddress.class);
     }
-
-    public void assumedMessageSize(int bytes) {
-      this.assumedMessageSize = bytes;
-    }
   }
 }

inprocess/src/main/java/io/grpc/inprocess/InProcessServer.java

@@ -62,7 +62,7 @@ static InProcessServer findServer(SocketAddress addr) {
    */
   private ScheduledExecutorService scheduler;
 
-  private long assumedMessageSize = -1;
+  private final long assumedMessageSize;
 
   InProcessServer(
       InProcessServerBuilder builder,

inprocess/src/main/java/io/grpc/inprocess/InProcessServerBuilder.java

@@ -212,8 +212,8 @@ void setStatsEnabled(boolean value) {
     this.serverImplBuilder.setStatsEnabled(value);
   }
 
-  public InProcessServerBuilder assumedMessageSize(long bytes) {
-    this.assumedMessageSize = bytes;
+  public InProcessServerBuilder assumedMessageSize(long assumedMessageSize) {
+    this.assumedMessageSize = assumedMessageSize;
     return this;
   }
 }

inprocess/src/main/java/io/grpc/inprocess/InProcessTransport.java

@@ -99,9 +99,8 @@ final class InProcessTransport implements ServerTransport, ConnectionClientTrans
   private ServerTransportListener serverTransportListener;
   private Attributes serverStreamAttributes;
   private ManagedClientTransport.Listener clientTransportListener;
-  // These sizes are assumed from the sender' side.
-  private long serverAssumedMessageSize = -1;
-   private long clientAssumedMessageSize = -1;
+  // The size is assumed from the sender's side.
+  private long assumedMessageSize = -1;
   @GuardedBy("this")
   private boolean shutdown;
   @GuardedBy("this")
@@ -172,7 +171,7 @@ public InProcessTransport(
           Attributes eagAttrs, boolean includeCauseWithStatus, long assumedMessageSize) {
     this(address, maxInboundMetadataSize, authority, userAgent, eagAttrs,
             Optional.<ServerListener>absent(), includeCauseWithStatus);
-            this.clientAssumedMessageSize = assumedMessageSize;
+    this.assumedMessageSize = assumedMessageSize;
   }
 
   InProcessTransport(
@@ -201,9 +200,9 @@ public synchronized Runnable start(ManagedClientTransport.Listener listener) {
         serverSchedulerPool = server.getScheduledExecutorServicePool();
         serverScheduler = serverSchedulerPool.getObject();
         serverStreamTracerFactories = server.getStreamTracerFactories();
+        assumedMessageSize = server.getAssumedMessageSize();
         // Must be semi-initialized; past this point, can begin receiving requests
         serverTransportListener = server.register(this);
-        serverAssumedMessageSize = server.getAssumedMessageSize();
       }
     }
     if (serverTransportListener == null) {
@@ -430,7 +429,7 @@ private void streamClosed() {
       }
     }
 
-    private long getMessageLength(InputStream inputStream) {
+    private long getKnownLength(InputStream inputStream) {
       try {
         return inputStream.available();
       } catch (IOException | RuntimeException e) {
@@ -543,22 +542,21 @@ public void writeMessage(InputStream message) {
           clientStream.statsTraceCtx.inboundMessage(outboundSeqNo);
           clientStream.statsTraceCtx.inboundMessageRead(outboundSeqNo, -1, -1);
 
-          /*long messageLength = serverAssumedMessageSize == -1
-                  ? getMessageLength(message) : serverAssumedMessageSize;*/
-          long messageLength = -1;
+          long messageLength;
           try {
-            if (!(message instanceof KnownLength) && !(message instanceof ByteArrayInputStream)) {
+            if (assumedMessageSize != -1) {
+              messageLength = assumedMessageSize;
+            } else if (message instanceof KnownLength || message instanceof ByteArrayInputStream) {
+              messageLength = getKnownLength(message);
+            } else {
               InputStream oldMessage = message;
               byte[] payload = ByteStreams.toByteArray(message);
               messageLength = payload.length;
               message = new ByteArrayInputStream(payload);
               oldMessage.close();
             }
-            else {
-              messageLength = getMessageLength(message);
-            }
-          } catch (IOException ex) {
-            throw new RuntimeException(ex);
+          } catch (IOException e) {
+            throw new RuntimeException("Error processing the message length", e);
           }
 
           statsTraceCtx.outboundUncompressedSize(messageLength);
@@ -838,22 +836,21 @@ public void writeMessage(InputStream message) {
           serverStream.statsTraceCtx.inboundMessage(outboundSeqNo);
           serverStream.statsTraceCtx.inboundMessageRead(outboundSeqNo, -1, -1);
 
-          /*long messageLength = clientAssumedMessageSize == -1
-                  ? getMessageLength(message) : clientAssumedMessageSize;*/
-          long messageLength = -1;
+          long messageLength;
           try {
-            if (!(message instanceof KnownLength) && !(message instanceof ByteArrayInputStream)) {
+            if (assumedMessageSize != -1) {
+              messageLength = assumedMessageSize;
+            } else if (message instanceof KnownLength || message instanceof ByteArrayInputStream) {
+              messageLength = getKnownLength(message);
+            } else {
               InputStream oldMessage = message;
               byte[] payload = ByteStreams.toByteArray(message);
               messageLength = payload.length;
               message = new ByteArrayInputStream(payload);
               oldMessage.close();
             }
-            else {
-              messageLength = getMessageLength(message);
-            }
-          } catch (IOException ex) {
-            throw new RuntimeException(ex);
+          } catch (IOException e) {
+            throw new RuntimeException("Error processing the message length", e);
           }
           statsTraceCtx.outboundUncompressedSize(messageLength);
           statsTraceCtx.outboundWireSize(messageLength);

inprocess/src/test/java/io/grpc/inprocess/AnonymousInProcessTransportTest.java

@@ -52,6 +52,6 @@ protected InternalServer newServer(
   protected ManagedClientTransport newClientTransport(InternalServer server) {
     return new InProcessTransport(
         address, GrpcUtil.DEFAULT_MAX_HEADER_LIST_SIZE,
-        testAuthority(server), USER_AGENT, eagAttrs(), false, -1);
+        testAuthority(server), USER_AGENT, eagAttrs(), false);
   }
 }

inprocess/src/test/java/io/grpc/inprocess/InProcessTransportTest.java

@@ -83,7 +83,7 @@ protected String testAuthority(InternalServer server) {
   protected ManagedClientTransport newClientTransport(InternalServer server) {
     return new InProcessTransport(
         new InProcessSocketAddress(TRANSPORT_NAME), GrpcUtil.DEFAULT_MAX_HEADER_LIST_SIZE,
-        testAuthority(server), USER_AGENT, eagAttrs(), false, -1);
+        testAuthority(server), USER_AGENT, eagAttrs(), false);
   }
 
   @Test

inprocess/src/test/java/io/grpc/inprocess/SizesReportedInProcessTransportTest.java

@@ -0,0 +1,68 @@
+/*
+ * Copyright 2021 The gRPC Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.grpc.inprocess;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import io.grpc.ServerStreamTracer;
+import io.grpc.internal.GrpcUtil;
+import io.grpc.internal.InternalServer;
+import io.grpc.internal.ManagedClientTransport;
+import io.grpc.internal.testing.TestStreamTracer;
+import java.util.List;
+import org.junit.After;
+import org.junit.Assert;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class SizesReportedInProcessTransportTest extends InProcessTransportTest {
+  private static final long TEST_MESSAGE_LENGTH = 100;
+  private AnonymousInProcessSocketAddress address = new AnonymousInProcessSocketAddress();
+
+  @After
+  @Override
+  public void tearDown() throws InterruptedException {
+    super.tearDown();
+    assertThat(address.getServer()).isNull();
+  }
+
+  @Override
+  protected InternalServer newServer(
+      List<ServerStreamTracer.Factory> streamTracerFactories) {
+    InProcessServerBuilder builder = InProcessServerBuilder.forAddress(address)
+              .maxInboundMetadataSize(GrpcUtil.DEFAULT_MAX_HEADER_LIST_SIZE)
+            .assumedMessageSize(TEST_MESSAGE_LENGTH);
+    return new InProcessServer(builder, streamTracerFactories);
+  }
+
+  @Override
+  protected ManagedClientTransport newClientTransport(InternalServer server) {
+    return new InProcessTransport(
+            address, GrpcUtil.DEFAULT_MAX_HEADER_LIST_SIZE,
+            testAuthority(server), USER_AGENT, eagAttrs(), false, TEST_MESSAGE_LENGTH);
+  }
+
+  @Override
+  public void assertInProcessTransportAssumedMessageSize(
+          TestStreamTracer streamTracerSender, TestStreamTracer streamTracerReceiver) {
+    Assert.assertEquals(TEST_MESSAGE_LENGTH, streamTracerSender.getOutboundWireSize());
+    Assert.assertEquals(TEST_MESSAGE_LENGTH, streamTracerSender.getOutboundUncompressedSize());
+    Assert.assertEquals(TEST_MESSAGE_LENGTH, streamTracerReceiver.getInboundWireSize());
+    Assert.assertEquals(TEST_MESSAGE_LENGTH, streamTracerReceiver.getInboundUncompressedSize());
+  }
+}

opentelemetry/src/test/java/io/grpc/opentelemetry/OpenTelemetryMetricsModuleTest.java

@@ -37,6 +37,7 @@
 import io.grpc.ClientInterceptor;
 import io.grpc.ClientInterceptors;
 import io.grpc.ClientStreamTracer;
+import io.grpc.KnownLength;
 import io.grpc.Metadata;
 import io.grpc.MethodDescriptor;
 import io.grpc.ServerCall;
@@ -53,6 +54,7 @@
 import io.opentelemetry.api.metrics.Meter;
 import io.opentelemetry.sdk.common.InstrumentationScopeInfo;
 import io.opentelemetry.sdk.testing.junit4.OpenTelemetryRule;
+import java.io.IOException;
 import java.io.InputStream;
 import java.util.Arrays;
 import java.util.Map;
@@ -107,7 +109,7 @@ public class OpenTelemetryMetricsModuleTest {
       { 0L, 1024L, 2048L, 4096L, 16384L, 65536L, 262144L, 1048576L, 4194304L, 16777216L,
       67108864L, 268435456L, 1073741824L, 4294967296L };
 
-  private static final class StringInputStream extends InputStream {
+  private static final class StringInputStream extends InputStream implements KnownLength {
     final String string;
 
     StringInputStream(String string) {
@@ -118,6 +120,11 @@ private static final class StringInputStream extends InputStream {
     public int read() {
       throw new UnsupportedOperationException("should not be called");
     }
+
+    @Override
+    public int available() throws IOException {
+      return string == null ? 0 : string.length();
+    }
   }
 
   private static final MethodDescriptor.Marshaller<String> MARSHALLER =