[#691] Improve memory usage via reuse InputStreamBufferInput

koo-taejin · koo-taejin · commit 56ac516c03b3 · 2023-01-19T16:08:30.000+09:00
diff --git a/msgpack-core/src/main/java/org/msgpack/core/buffer/ArrayBufferInput.java b/msgpack-core/src/main/java/org/msgpack/core/buffer/ArrayBufferInput.java
@@ -21,7 +21,7 @@
  * MessageBufferInput adapter for byte arrays
  */
 public class ArrayBufferInput
-        implements MessageBufferInput
+        implements MessageBufferInput<byte[]>
 {
     private MessageBuffer buffer;
     private boolean isEmpty;
@@ -66,9 +66,14 @@ public MessageBuffer reset(MessageBuffer buf)
         return old;
     }
 
-    public void reset(byte[] arr)
+    @Override
+    public byte[] reset(byte[] arr)
     {
-        reset(MessageBuffer.wrap(checkNotNull(arr, "input array is null")));
+        final MessageBuffer messageBuffer = reset(MessageBuffer.wrap(checkNotNull(arr, "input array is null")));
+        if (messageBuffer == null) {
+            return null;
+        }
+        return messageBuffer.array();
     }
 
     public void reset(byte[] arr, int offset, int len)
diff --git a/msgpack-core/src/main/java/org/msgpack/core/buffer/ByteBufferInput.java b/msgpack-core/src/main/java/org/msgpack/core/buffer/ByteBufferInput.java
@@ -23,7 +23,7 @@
  * {@link MessageBufferInput} adapter for {@link java.nio.ByteBuffer}
  */
 public class ByteBufferInput
-        implements MessageBufferInput
+        implements MessageBufferInput<ByteBuffer>
 {
     private ByteBuffer input;
     private boolean isRead = false;
@@ -39,6 +39,7 @@ public ByteBufferInput(ByteBuffer input)
      * @param input new buffer
      * @return the old buffer
      */
+    @Override
     public ByteBuffer reset(ByteBuffer input)
     {
         ByteBuffer old = this.input;
diff --git a/msgpack-core/src/main/java/org/msgpack/core/buffer/ChannelBufferInput.java b/msgpack-core/src/main/java/org/msgpack/core/buffer/ChannelBufferInput.java
@@ -26,7 +26,7 @@
  * {@link MessageBufferInput} adapter for {@link java.nio.channels.ReadableByteChannel}
  */
 public class ChannelBufferInput
-        implements MessageBufferInput
+        implements MessageBufferInput<ReadableByteChannel>
 {
     private ReadableByteChannel channel;
     private final MessageBuffer buffer;
@@ -49,8 +49,8 @@ public ChannelBufferInput(ReadableByteChannel channel, int bufferSize)
      * @param channel new channel
      * @return the old resource
      */
+    @Override
     public ReadableByteChannel reset(ReadableByteChannel channel)
-            throws IOException
     {
         ReadableByteChannel old = this.channel;
         this.channel = channel;
diff --git a/msgpack-core/src/main/java/org/msgpack/core/buffer/InputStreamBufferInput.java b/msgpack-core/src/main/java/org/msgpack/core/buffer/InputStreamBufferInput.java
@@ -26,7 +26,7 @@
  * {@link MessageBufferInput} adapter for {@link InputStream}
  */
 public class InputStreamBufferInput
-        implements MessageBufferInput
+        implements MessageBufferInput<InputStream>
 {
     private InputStream in;
     private final byte[] buffer;
@@ -60,8 +60,8 @@ public InputStreamBufferInput(InputStream in, int bufferSize)
      * @param in new stream
      * @return the old resource
      */
+    @Override
     public InputStream reset(InputStream in)
-            throws IOException
     {
         InputStream old = this.in;
         this.in = in;
diff --git a/msgpack-core/src/main/java/org/msgpack/core/buffer/MessageBufferInput.java b/msgpack-core/src/main/java/org/msgpack/core/buffer/MessageBufferInput.java
@@ -24,7 +24,7 @@
  * A MessageBufferInput implementation has control of lifecycle of the memory so that it can reuse previously
  * allocated memory, use memory pools, or use memory-mapped files.
  */
-public interface MessageBufferInput
+public interface MessageBufferInput<T>
         extends Closeable
 {
     /**
@@ -40,6 +40,8 @@ public interface MessageBufferInput
     MessageBuffer next()
             throws IOException;
 
+    T reset(T input);
+
     /**
      * Closes the input.
      * <p>
diff --git a/msgpack-core/src/test/java/org/msgpack/core/buffer/SequenceMessageBufferInput.java b/msgpack-core/src/test/java/org/msgpack/core/buffer/SequenceMessageBufferInput.java
@@ -24,7 +24,7 @@
  * {@link MessageBufferInput} adapter for {@link MessageBufferInput} Enumeration
  */
 public class SequenceMessageBufferInput
-        implements MessageBufferInput
+        implements MessageBufferInput<Void>
 {
     private Enumeration<? extends MessageBufferInput> sequence;
     private MessageBufferInput input;
@@ -54,6 +54,11 @@ public MessageBuffer next() throws IOException
         return buffer;
     }
 
+    @Override
+    public Void reset(Void input) {
+        throw new UnsupportedOperationException("reset");
+    }
+
     private void nextInput() throws IOException
     {
         if (input != null) {
diff --git a/msgpack-core/src/test/scala/org/msgpack/core/MessageUnpackerTest.scala b/msgpack-core/src/test/scala/org/msgpack/core/MessageUnpackerTest.scala
@@ -29,7 +29,7 @@ import scala.jdk.CollectionConverters._
 import scala.util.Random
 
 object MessageUnpackerTest {
-  class SplitMessageBufferInput(array: Array[Array[Byte]]) extends MessageBufferInput {
+  class SplitMessageBufferInput(array: Array[Array[Byte]]) extends MessageBufferInput[Void] {
     var cursor = 0
     override def next(): MessageBuffer = {
       if (cursor < array.length) {
@@ -41,6 +41,9 @@ object MessageUnpackerTest {
       }
     }
 
+
+    override def reset(input: Void): Void = throw new UnsupportedOperationException("reset")
+
     override def close(): Unit = {}
   }
 }
diff --git a/msgpack-core/src/test/scala/org/msgpack/core/buffer/ByteStringTest.scala b/msgpack-core/src/test/scala/org/msgpack/core/buffer/ByteStringTest.scala
@@ -26,7 +26,7 @@ class ByteStringTest extends AirSpec {
   private val byteString     = ByteString(createMessagePackData(_.packString(unpackedString)))
 
   private def unpackString(messageBuffer: MessageBuffer) = {
-    val input = new MessageBufferInput {
+    val input = new MessageBufferInput[Void] {
 
       private var isRead = false
 
@@ -38,6 +38,8 @@ class ByteStringTest extends AirSpec {
           messageBuffer
         }
       override def close(): Unit = {}
+
+      override def reset(input: Void): Void = throw new UnsupportedOperationException("reset")
     }
 
     MessagePack.newDefaultUnpacker(input).unpackString()
diff --git a/msgpack-jackson/src/main/java/org/msgpack/jackson/dataformat/MessageBufferInputLocator.java b/msgpack-jackson/src/main/java/org/msgpack/jackson/dataformat/MessageBufferInputLocator.java
@@ -0,0 +1,23 @@
+//
+// MessagePack for Java
+//
+//    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 org.msgpack.jackson.dataformat;
+
+import org.msgpack.core.buffer.MessageBufferInput;
+
+public interface MessageBufferInputLocator
+{
+    MessageBufferInput get(Class clazz);
+}
diff --git a/msgpack-jackson/src/main/java/org/msgpack/jackson/dataformat/MessageBufferInputProvider.java b/msgpack-jackson/src/main/java/org/msgpack/jackson/dataformat/MessageBufferInputProvider.java
@@ -0,0 +1,23 @@
+//
+// MessagePack for Java
+//
+//    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 org.msgpack.jackson.dataformat;
+
+import org.msgpack.core.buffer.MessageBufferInput;
+
+interface MessageBufferInputProvider
+{
+    MessageBufferInput provide();
+}
diff --git a/msgpack-jackson/src/main/java/org/msgpack/jackson/dataformat/MessageBufferInputRegistry.java b/msgpack-jackson/src/main/java/org/msgpack/jackson/dataformat/MessageBufferInputRegistry.java
@@ -0,0 +1,46 @@
+//
+// MessagePack for Java
+//
+//    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 org.msgpack.jackson.dataformat;
+
+import org.msgpack.core.buffer.MessageBufferInput;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Objects;
+
+public class MessageBufferInputRegistry implements MessageBufferInputLocator
+{
+    private final Map<Class, MessageBufferInput> messageBufferInputMap = new HashMap<>(1);
+
+    @Override
+    public MessageBufferInput get(Class clazz)
+    {
+        return messageBufferInputMap.get(clazz);
+    }
+
+    public boolean register(Class clazz, MessageBufferInputProvider provider)
+    {
+        Objects.requireNonNull(clazz, "clazz");
+        Objects.requireNonNull(provider, "provider");
+
+        if (messageBufferInputMap.containsKey(clazz)) {
+            return false;
+        }
+
+        messageBufferInputMap.put(clazz, provider.provide());
+        return true;
+    }
+}
diff --git a/msgpack-jackson/src/main/java/org/msgpack/jackson/dataformat/MessagePackParser.java b/msgpack-jackson/src/main/java/org/msgpack/jackson/dataformat/MessagePackParser.java
@@ -47,8 +47,8 @@
 public class MessagePackParser
         extends ParserMinimalBase
 {
-    private static final ThreadLocal<Tuple<Object, MessageUnpacker>> messageUnpackerHolder =
-            new ThreadLocal<Tuple<Object, MessageUnpacker>>();
+    private static final ThreadLocal<Triple<Object, MessageUnpacker, MessageBufferInputLocator>> reuseObjectHolder =
+            new ThreadLocal<>();
     private final MessageUnpacker messageUnpacker;
 
     private static final BigInteger LONG_MIN = BigInteger.valueOf((long) Long.MIN_VALUE);
@@ -126,11 +126,17 @@ public MessagePackParser(
             IOContext ctxt,
             int features,
             ObjectCodec objectCodec,
-            InputStream in,
+            final InputStream in,
             boolean reuseResourceInParser)
             throws IOException
     {
-        this(ctxt, features, new InputStreamBufferInput(in), objectCodec, in, reuseResourceInParser);
+        this(ctxt, features, new MessageBufferInputProvider() {
+            @Override
+            public MessageBufferInput provide()
+            {
+                return new InputStreamBufferInput(in);
+            }
+        }, objectCodec, in, reuseResourceInParser);
     }
 
     public MessagePackParser(IOContext ctxt, int features, ObjectCodec objectCodec, byte[] bytes)
@@ -143,16 +149,22 @@ public MessagePackParser(
             IOContext ctxt,
             int features,
             ObjectCodec objectCodec,
-            byte[] bytes,
+            final byte[] bytes,
             boolean reuseResourceInParser)
             throws IOException
     {
-        this(ctxt, features, new ArrayBufferInput(bytes), objectCodec, bytes, reuseResourceInParser);
+        this(ctxt, features, new MessageBufferInputProvider() {
+            @Override
+            public MessageBufferInput provide()
+            {
+                return new ArrayBufferInput(bytes);
+            }
+        }, objectCodec, bytes, reuseResourceInParser);
     }
 
     private MessagePackParser(IOContext ctxt,
             int features,
-            MessageBufferInput input,
+            MessageBufferInputProvider bufferInputProvider,
             ObjectCodec objectCodec,
             Object src,
             boolean reuseResourceInParser)
@@ -167,29 +179,48 @@ private MessagePackParser(IOContext ctxt,
         parsingContext = JsonReadContext.createRootContext(dups);
         this.reuseResourceInParser = reuseResourceInParser;
         if (!reuseResourceInParser) {
-            this.messageUnpacker = MessagePack.newDefaultUnpacker(input);
+            this.messageUnpacker = MessagePack.newDefaultUnpacker(bufferInputProvider.provide());
             return;
         }
         else {
             this.messageUnpacker = null;
         }
 
         MessageUnpacker messageUnpacker;
-        Tuple<Object, MessageUnpacker> messageUnpackerTuple = messageUnpackerHolder.get();
-        if (messageUnpackerTuple == null) {
-            messageUnpacker = MessagePack.newDefaultUnpacker(input);
+        MessageBufferInputLocator messageBufferInputLocator;
+        Triple<Object, MessageUnpacker, MessageBufferInputLocator> messageUnpackerTriple = reuseObjectHolder.get();
+        if (messageUnpackerTriple == null) {
+            final MessageBufferInputRegistry messageBufferInputRegistry = new MessageBufferInputRegistry();
+            messageBufferInputRegistry.register(src.getClass(), bufferInputProvider);
+            messageBufferInputLocator = messageBufferInputRegistry;
+            messageUnpacker = MessagePack.newDefaultUnpacker(messageBufferInputRegistry.get(src.getClass()));
         }
         else {
             // Considering to reuse InputStream with JsonParser.Feature.AUTO_CLOSE_SOURCE,
             // MessagePackParser needs to use the MessageUnpacker that has the same InputStream
             // since it has buffer which has loaded the InputStream data ahead.
             // However, it needs to call MessageUnpacker#reset when the source is different from the previous one.
-            if (isEnabled(JsonParser.Feature.AUTO_CLOSE_SOURCE) || messageUnpackerTuple.first() != src) {
-                messageUnpackerTuple.second().reset(input);
+            if (isEnabled(JsonParser.Feature.AUTO_CLOSE_SOURCE) || messageUnpackerTriple.first() != src) {
+                final MessageBufferInputLocator bufferInputLocator = messageUnpackerTriple.third();
+                MessageBufferInput messageBufferInput = bufferInputLocator.get(src.getClass());
+                if (messageBufferInput != null) {
+                    messageBufferInput.reset(src);
+                }
+                else {
+                    if (bufferInputLocator instanceof MessageBufferInputRegistry) {
+                        ((MessageBufferInputRegistry) bufferInputLocator).register(src.getClass(), bufferInputProvider);
+                        messageBufferInput = bufferInputLocator.get(src.getClass());
+                    }
+                    else {
+                        messageBufferInput = bufferInputProvider.provide();
+                    }
+                }
+                messageUnpackerTriple.second().reset(messageBufferInput);
             }
-            messageUnpacker = messageUnpackerTuple.second();
+            messageUnpacker = messageUnpackerTriple.second();
+            messageBufferInputLocator = messageUnpackerTriple.third();
         }
-        messageUnpackerHolder.set(new Tuple<Object, MessageUnpacker>(src, messageUnpacker));
+        reuseObjectHolder.set(new Triple<Object, MessageUnpacker, MessageBufferInputLocator>(src, messageUnpacker, messageBufferInputLocator));
     }
 
     public void setExtensionTypeCustomDeserializers(ExtensionTypeCustomDeserializers extTypeCustomDesers)
@@ -690,7 +721,7 @@ private MessageUnpacker getMessageUnpacker()
             return this.messageUnpacker;
         }
 
-        Tuple<Object, MessageUnpacker> messageUnpackerTuple = messageUnpackerHolder.get();
+        Triple<Object, MessageUnpacker, MessageBufferInputLocator> messageUnpackerTuple = reuseObjectHolder.get();
         if (messageUnpackerTuple == null) {
             throw new IllegalStateException("messageUnpacker is null");
         }
diff --git a/msgpack-jackson/src/main/java/org/msgpack/jackson/dataformat/Triple.java b/msgpack-jackson/src/main/java/org/msgpack/jackson/dataformat/Triple.java

Original file line number	Diff line number	Diff line change
`@@ -21,7 +21,7 @@`
`21`	`21`	`* MessageBufferInput adapter for byte arrays`
`22`	`22`	`*/`
`23`	`23`	`public class ArrayBufferInput`
`24`		`- implements MessageBufferInput`
	`24`	`+ implements MessageBufferInput<byte[]>`
`25`	`25`	`{`
`26`	`26`	`private MessageBuffer buffer;`
`27`	`27`	`private boolean isEmpty;`
`@@ -66,9 +66,14 @@ public MessageBuffer reset(MessageBuffer buf)`
`66`	`66`	`return old;`
`67`	`67`	`}`
`68`	`68`
`69`		`- public void reset(byte[] arr)`
	`69`	`+ @Override`
	`70`	`+ public byte[] reset(byte[] arr)`
`70`	`71`	`{`
`71`		`- reset(MessageBuffer.wrap(checkNotNull(arr, "input array is null")));`
	`72`	`+ final MessageBuffer messageBuffer = reset(MessageBuffer.wrap(checkNotNull(arr, "input array is null")));`
	`73`	`+ if (messageBuffer == null) {`
	`74`	`+ return null;`
	`75`	`+ }`
	`76`	`+ return messageBuffer.array();`
`72`	`77`	`}`
`73`	`78`
`74`	`79`	`public void reset(byte[] arr, int offset, int len)`
Original file line number	Diff line number	Diff line change
`@@ -24,7 +24,7 @@`
`24`	`24`	`* {@link MessageBufferInput} adapter for {@link MessageBufferInput} Enumeration`
`25`	`25`	`*/`
`26`	`26`	`public class SequenceMessageBufferInput`
`27`		`- implements MessageBufferInput`
	`27`	`+ implements MessageBufferInput<Void>`
`28`	`28`	`{`
`29`	`29`	`private Enumeration<? extends MessageBufferInput> sequence;`
`30`	`30`	`private MessageBufferInput input;`
`@@ -54,6 +54,11 @@ public MessageBuffer next() throws IOException`
`54`	`54`	`return buffer;`
`55`	`55`	`}`
`56`	`56`
	`57`	`+ @Override`
	`58`	`+ public Void reset(Void input) {`
	`59`	`+ throw new UnsupportedOperationException("reset");`
	`60`	`+ }`
	`61`	`+`
`57`	`62`	`private void nextInput() throws IOException`
`58`	`63`	`{`
`59`	`64`	`if (input != null) {`
Original file line number	Diff line number	Diff line change
`@@ -29,7 +29,7 @@ import scala.jdk.CollectionConverters._`
`29`	`29`	`import scala.util.Random`
`30`	`30`
`31`	`31`	`object MessageUnpackerTest {`
`32`		`- class SplitMessageBufferInput(array: Array[Array[Byte]]) extends MessageBufferInput {`
	`32`	`+ class SplitMessageBufferInput(array: Array[Array[Byte]]) extends MessageBufferInput[Void] {`
`33`	`33`	`var cursor = 0`
`34`	`34`	`override def next(): MessageBuffer = {`
`35`	`35`	`if (cursor < array.length) {`
`@@ -41,6 +41,9 @@ object MessageUnpackerTest {`
`41`	`41`	`}`
`42`	`42`	`}`
`43`	`43`
	`44`	`+`
	`45`	`+ override def reset(input: Void): Void = throw new UnsupportedOperationException("reset")`
	`46`	`+`
`44`	`47`	`override def close(): Unit = {}`
`45`	`48`	`}`
`46`	`49`	`}`
Original file line number	Diff line number	Diff line change
`@@ -26,7 +26,7 @@ class ByteStringTest extends AirSpec {`
`26`	`26`	`private val byteString = ByteString(createMessagePackData(_.packString(unpackedString)))`
`27`	`27`
`28`	`28`	`private def unpackString(messageBuffer: MessageBuffer) = {`
`29`		`- val input = new MessageBufferInput {`
	`29`	`+ val input = new MessageBufferInput[Void] {`
`30`	`30`
`31`	`31`	`private var isRead = false`
`32`	`32`
`@@ -38,6 +38,8 @@ class ByteStringTest extends AirSpec {`
`38`	`38`	`messageBuffer`
`39`	`39`	`}`
`40`	`40`	`override def close(): Unit = {}`
	`41`	`+`
	`42`	`+ override def reset(input: Void): Void = throw new UnsupportedOperationException("reset")`
`41`	`43`	`}`
`42`	`44`
`43`	`45`	`MessagePack.newDefaultUnpacker(input).unpackString()`