Add ByteBuffer input support again

Author: miniway

msgpack-core/src/main/java/org/msgpack/core/MessagePack.java

@@ -16,6 +16,7 @@
 package org.msgpack.core;
 
 import org.msgpack.core.buffer.ArrayBufferInput;
+import org.msgpack.core.buffer.ByteBufferInput;
 import org.msgpack.core.buffer.ChannelBufferInput;
 import org.msgpack.core.buffer.ChannelBufferOutput;
 import org.msgpack.core.buffer.InputStreamBufferInput;
@@ -25,6 +26,7 @@
 
 import java.io.InputStream;
 import java.io.OutputStream;
+import java.nio.ByteBuffer;
 import java.nio.channels.ReadableByteChannel;
 import java.nio.channels.WritableByteChannel;
 import java.nio.charset.Charset;
@@ -236,6 +238,17 @@ public static MessageUnpacker newDefaultUnpacker(byte[] contents, int offset, in
         return DEFAULT_UNPACKER_CONFIG.newUnpacker(contents, offset, length);
     }
 
+    /**
+     * Create an unpacker that reads the data from a given ByteBuffer
+     *
+     * @param contents
+     * @return
+     */
+    public static MessageUnpacker newDefaultUnpacker(ByteBuffer contents)
+    {
+        return DEFAULT_UNPACKER_CONFIG.newUnpacker(contents);
+    }
+
     /**
      * MessagePacker configuration.
      */
@@ -524,6 +537,17 @@ public MessageUnpacker newUnpacker(byte[] contents, int offset, int length)
             return newUnpacker(new ArrayBufferInput(contents, offset, length));
         }
 
+        /**
+         * Create an unpacker that reads the data from a given ByteBuffer
+         *
+         * @param contents
+         * @return
+         */
+        public MessageUnpacker newUnpacker(ByteBuffer contents)
+        {
+            return newUnpacker(new ByteBufferInput(contents));
+        }
+
         /**
          * Allow unpackBinaryHeader to read str format family  (default: true)
          */

msgpack-core/src/main/java/org/msgpack/core/MessagePacker.java

@@ -531,8 +531,7 @@ else if (s.length() < (1 << 8)) {
                         throw new IllegalArgumentException("Unexpected UTF-8 encoder state");
                     }
                     // move 1 byte backward to expand 3-byte header region to 3 bytes
-                    buffer.putBytes(position + 3,
-                            buffer.array(), buffer.arrayOffset() + position + 2, written);
+                    buffer.putMessageBuffer(position + 3, buffer, position + 2, written);
                     // write 3-byte header
                     buffer.putByte(position++, STR16);
                     buffer.putShort(position, (short) written);
@@ -560,8 +559,7 @@ else if (s.length() < (1 << 16)) {
                         throw new IllegalArgumentException("Unexpected UTF-8 encoder state");
                     }
                     // move 2 bytes backward to expand 3-byte header region to 5 bytes
-                    buffer.putBytes(position + 5,
-                            buffer.array(), buffer.arrayOffset() + position + 3, written);
+                    buffer.putMessageBuffer(position + 5, buffer, position + 3, written);
                     // write 3-byte header header
                     buffer.putByte(position++, STR32);
                     buffer.putInt(position, written);

msgpack-core/src/main/java/org/msgpack/core/MessageUnpacker.java

@@ -213,14 +213,9 @@ private MessageBuffer prepareNumberBuffer(int readLength)
             // fill the temporary buffer from the current data fragment and
             // next fragment(s).
 
-            // TODO buffer.array() doesn't work if MessageBuffer is allocated by
-            //      newDirectBuffer. dd copy method to MessageBuffer to solve this issue.
-
             int off = 0;
             if (remaining > 0) {
-                numberBuffer.putBytes(0,
-                        buffer.array(), buffer.arrayOffset() + position,
-                        remaining);
+                numberBuffer.putMessageBuffer(0, buffer, position, remaining);
                 readLength -= remaining;
                 off += remaining;
             }
@@ -229,16 +224,12 @@ private MessageBuffer prepareNumberBuffer(int readLength)
                 nextBuffer();
                 int nextSize = buffer.size();
                 if (nextSize >= readLength) {
-                    numberBuffer.putBytes(off,
-                            buffer.array(), buffer.arrayOffset(),
-                            readLength);
+                    numberBuffer.putMessageBuffer(off, buffer, 0, readLength);
                     position = readLength;
                     break;
                 }
                 else {
-                    numberBuffer.putBytes(off,
-                            buffer.array(), buffer.arrayOffset(),
-                            nextSize);
+                    numberBuffer.putMessageBuffer(off, buffer, 0, nextSize);
                     readLength -= nextSize;
                     off += nextSize;
                 }
@@ -1041,7 +1032,8 @@ private void handleCoderError(CoderResult cr)
     private String decodeStringFastPath(int length)
     {
         if (actionOnMalformedString == CodingErrorAction.REPLACE &&
-                actionOnUnmappableString == CodingErrorAction.REPLACE) {
+                actionOnUnmappableString == CodingErrorAction.REPLACE &&
+                buffer.hasArray()) {
             String s = new String(buffer.array(), buffer.arrayOffset() + position, length, MessagePack.UTF8);
             position += length;
             return s;

msgpack-core/src/main/java/org/msgpack/core/buffer/ByteBufferInput.java

@@ -0,0 +1,69 @@
+//
+// 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.core.buffer;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+
+import static org.msgpack.core.Preconditions.checkNotNull;
+
+/**
+ * {@link MessageBufferInput} adapter for {@link java.nio.ByteBuffer}
+ */
+public class ByteBufferInput
+        implements MessageBufferInput
+{
+    private ByteBuffer input;
+    private boolean isRead = false;
+
+    public ByteBufferInput(ByteBuffer input)
+    {
+        this.input = checkNotNull(input, "input ByteBuffer is null");
+    }
+
+    /**
+     * Reset buffer. This method doesn't close the old resource.
+     *
+     * @param input new buffer
+     * @return the old resource
+     */
+    public ByteBuffer reset(ByteBuffer input)
+    {
+        ByteBuffer old = this.input;
+        this.input = input;
+        isRead = false;
+        return old;
+    }
+
+    @Override
+    public MessageBuffer next()
+            throws IOException
+    {
+        if (isRead) {
+            return null;
+        }
+
+        isRead = true;
+        return MessageBuffer.wrap(input);
+    }
+
+    @Override
+    public void close()
+            throws IOException
+    {
+        // Nothing to do
+    }
+}

msgpack-core/src/main/java/org/msgpack/core/buffer/MessageBuffer.java

@@ -43,6 +43,7 @@ public class MessageBuffer
      * Reference to MessageBuffer Constructors
      */
     private static final Constructor<?> mbArrConstructor;
+    private static final Constructor<?> mbBBConstructor;
 
     /**
      * The offset from the object memory header to its byte array data
@@ -145,6 +146,11 @@ public class MessageBuffer
                 Constructor<?> mbArrCstr = bufferCls.getDeclaredConstructor(byte[].class, int.class, int.class);
                 mbArrCstr.setAccessible(true);
                 mbArrConstructor = mbArrCstr;
+
+                // MessageBufferX(ByteBuffer) constructor
+                Constructor<?> mbBBCstr = bufferCls.getDeclaredConstructor(ByteBuffer.class);
+                mbBBCstr.setAccessible(true);
+                mbBBConstructor = mbBBCstr;
             }
             catch (Exception e) {
                 e.printStackTrace(System.err);
@@ -170,6 +176,12 @@ public class MessageBuffer
      */
     protected final int size;
 
+    /**
+     * Reference is used to hold a reference to an object that holds the underlying memory so that it cannot be
+     * released by the garbage collector.
+     */
+    protected final ByteBuffer reference;
+
     public static MessageBuffer allocate(int length)
     {
         return wrap(new byte[length]);
@@ -185,6 +197,11 @@ public static MessageBuffer wrap(byte[] array, int offset, int length)
         return newMessageBuffer(array, offset, length);
     }
 
+    public static MessageBuffer wrap(ByteBuffer bb)
+    {
+        return newMessageBuffer(bb).slice(bb.position(), bb.remaining());
+    }
+
     /**
      * Creates a new MessageBuffer instance backed by a java heap array
      *
@@ -202,11 +219,32 @@ private static MessageBuffer newMessageBuffer(byte[] arr, int off, int len)
         }
     }
 
+    /**
+     * Creates a new MessageBuffer instance backed by ByteBuffer
+     *
+     * @param bb
+     * @return
+     */
+    private static MessageBuffer newMessageBuffer(ByteBuffer bb)
+    {
+        checkNotNull(bb);
+        try {
+            // We need to use reflection to create MessageBuffer instances in order to prevent TypeProfile generation for getInt method. TypeProfile will be
+            // generated to resolve one of the method references when two or more classes overrides the method.
+            return (MessageBuffer) mbBBConstructor.newInstance(bb);
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
+
     public static void releaseBuffer(MessageBuffer buffer)
     {
         if (isUniversalBuffer || buffer.base instanceof byte[]) {
             // We have nothing to do. Wait until the garbage-collector collects this array object
         }
+        else if (DirectBufferAccess.isDirectByteBufferInstance(buffer.base)) {
+            DirectBufferAccess.clean(buffer.base);
+        }
         else {
             // Maybe cannot reach here
             unsafe.freeMemory(buffer.address);
@@ -225,13 +263,43 @@ public static void releaseBuffer(MessageBuffer buffer)
         this.base = arr;
         this.address = ARRAY_BYTE_BASE_OFFSET + offset;
         this.size = length;
+        this.reference = null;
+    }
+
+    /**
+     * Create a MessageBuffer instance from a given ByteBuffer instance
+     *
+     * @param bb
+     */
+    MessageBuffer(ByteBuffer bb)
+    {
+        if (bb.isDirect()) {
+            if (isUniversalBuffer) {
+                throw new IllegalStateException("Cannot create MessageBuffer from DirectBuffer");
+            }
+            // Direct buffer or off-heap memory
+            this.base = null;
+            this.address = DirectBufferAccess.getAddress(bb);
+            this.size = bb.capacity();
+            this.reference = bb;
+        }
+        else if (bb.hasArray()) {
+            this.base = bb.array();
+            this.address = ARRAY_BYTE_BASE_OFFSET;
+            this.size = bb.array().length;
+            this.reference = null;
+        }
+        else {
+            throw new IllegalArgumentException("Only the array-backed ByteBuffer or DirectBuffer are supported");
+        }
     }
 
     protected MessageBuffer(Object base, long address, int length)
     {
         this.base = base;
         this.address = address;
         this.size = length;
+        this.reference = null;
     }
 
     /**
@@ -393,6 +461,11 @@ else if (src.hasArray()) {
         }
     }
 
+    public void putMessageBuffer(int index, MessageBuffer src, int srcOffset, int len)
+    {
+        unsafe.copyMemory(src.base, src.address + srcOffset, base, address + index, len);
+    }
+
     /**
      * Create a ByteBuffer view of the range [index, index+length) of this memory
      *
@@ -402,7 +475,13 @@ else if (src.hasArray()) {
      */
     public ByteBuffer sliceAsByteBuffer(int index, int length)
     {
-        return ByteBuffer.wrap((byte[]) base, (int) ((address - ARRAY_BYTE_BASE_OFFSET) + index), length);
+        if (hasArray()) {
+            return ByteBuffer.wrap((byte[]) base, (int) ((address - ARRAY_BYTE_BASE_OFFSET) + index), length);
+        }
+        else {
+            assert (!isUniversalBuffer);
+            return DirectBufferAccess.newByteBuffer(address, index, length, reference);
+        }
     }
 
     /**
@@ -415,6 +494,11 @@ public ByteBuffer sliceAsByteBuffer()
         return sliceAsByteBuffer(0, size());
     }
 
+    public boolean hasArray()
+    {
+        return base instanceof byte[];
+    }
+
     /**
      * Get a copy of this buffer
      *