diff --git a/protobuf-lite/src/main/java/io/grpc/protobuf/lite/ProtoLiteUtils.java b/protobuf-lite/src/main/java/io/grpc/protobuf/lite/ProtoLiteUtils.java index 8cf95115fd..05db39db7f 100644 --- a/protobuf-lite/src/main/java/io/grpc/protobuf/lite/ProtoLiteUtils.java +++ b/protobuf-lite/src/main/java/io/grpc/protobuf/lite/ProtoLiteUtils.java @@ -49,6 +49,8 @@ import io.grpc.internal.GrpcUtil; import java.io.IOException; import java.io.InputStream; +import java.lang.ref.Reference; +import java.lang.ref.WeakReference; /** * Utility methods for using protobuf with grpc. @@ -78,15 +80,11 @@ public class ProtoLiteUtils { globalRegistry = checkNotNull(newRegistry, "newRegistry"); } - /** - * Local cache of buffers to use for parsing. ThreadLocal used a WeakReference internally, so - * these will not be retained. - */ - private static final ThreadLocal bufs = new ThreadLocal() { + private static final ThreadLocal> bufs = new ThreadLocal>() { @Override - protected byte[] initialValue() { - return new byte[4096]; // Picked at random. + protected Reference initialValue() { + return new WeakReference(new byte[4096]); // Picked at random. } }; @@ -141,11 +139,11 @@ public class ProtoLiteUtils { if (stream instanceof KnownLength) { int size = stream.available(); if (size > 0 && size <= GrpcUtil.DEFAULT_MAX_MESSAGE_SIZE) { - // Coded Input stream does not escape, so buf does not escape. - byte[] buf = bufs.get(); - if (buf.length < size) { + // buf should not be used after this method has returned. + byte[] buf = bufs.get().get(); + if (buf == null || buf.length < size) { buf = new byte[size]; - bufs.set(buf); + bufs.set(new WeakReference(buf)); } int chunkSize; int position = 0;