Add tests

apache · Dec 27, 2023 · 038e832 · 038e832
1 parent 74277cf
commit 038e832
Show file tree

Hide file tree

Showing 2 changed files with 55 additions and 2 deletions.
diff --git a/remoting/src/main/java/org/apache/rocketmq/remoting/netty/FileRegionEncoder.java b/remoting/src/main/java/org/apache/rocketmq/remoting/netty/FileRegionEncoder.java
@@ -57,8 +57,9 @@ public int write(ByteBuffer src) {
                 // To prevent mem_copy.
                 CompositeByteBuf b = (CompositeByteBuf) out;
                 // Have to increase writerIndex manually.
-                b.addComponent(true, Unpooled.wrappedBuffer(src));
-                return out.capacity();
+                ByteBuf unpooled = Unpooled.wrappedBuffer(src);
+                b.addComponent(true, unpooled);
+                return unpooled.readableBytes();
             }
 
             @Override

diff --git a/remoting/src/test/java/org/apache/rocketmq/remoting/netty/FileRegionEncoderTest.java b/remoting/src/test/java/org/apache/rocketmq/remoting/netty/FileRegionEncoderTest.java
@@ -18,17 +18,22 @@
 package org.apache.rocketmq.remoting.netty;
 
 import io.netty.buffer.ByteBuf;
+import io.netty.buffer.ByteBufAllocator;
+import io.netty.buffer.CompositeByteBuf;
+import io.netty.buffer.Unpooled;
 import io.netty.channel.DefaultFileRegion;
 import io.netty.channel.FileRegion;
 import io.netty.channel.embedded.EmbeddedChannel;
 import java.io.BufferedOutputStream;
 import java.io.File;
 import java.io.FileOutputStream;
 import java.io.IOException;
+import java.nio.ByteBuffer;
 import java.util.Random;
 import java.util.UUID;
 import org.junit.Assert;
 import org.junit.Test;
+import sun.nio.ch.DirectBuffer;
 
 public class FileRegionEncoderTest {
 
@@ -77,4 +82,51 @@ private static void write(File file, byte[] data) throws IOException {
             }
         }
     }
+
+    @Test
+    public void testUnpooledByteBuffer() {
+        ByteBuffer buffer = ByteBuffer.allocateDirect(128);
+        try {
+            int value = Integer.MAX_VALUE;
+            for (int i = 0; i < 128 / 4; i++) {
+                buffer.putInt(value);
+            }
+            buffer.flip();
+            ByteBuf buf1 = Unpooled.wrappedBuffer(buffer);
+            Assert.assertEquals(buf1.readableBytes(), 128);
+            Assert.assertEquals(buf1.readerIndex(), 0);
+            Assert.assertEquals(buf1.writerIndex(), 128);
+
+            for (int i = 0; i < 128 / 4; i++) {
+                Assert.assertEquals(buf1.readInt(), value);
+            }
+        } finally {
+            ((DirectBuffer) buffer).cleaner().clean();
+        }
+    }
+
+    @Test
+    public void testCompositeByteBuffer() {
+        ByteBuffer buffer = ByteBuffer.allocateDirect(128);
+        try {
+            int value = Integer.MAX_VALUE;
+            for (int i = 0; i < 128 / 4; i++) {
+                buffer.putInt(value);
+            }
+            buffer.flip();
+            ByteBuf buf = Unpooled.wrappedBuffer(buffer);
+            CompositeByteBuf cbuf = ByteBufAllocator.DEFAULT.compositeBuffer();
+            cbuf.addComponent(true, buf);
+            Assert.assertEquals(cbuf.readerIndex(), 0);
+            Assert.assertEquals(cbuf.writerIndex(), 128);
+            Assert.assertEquals(cbuf.readableBytes(), 128);
+            Assert.assertEquals(cbuf.refCnt(), 1);
+            Assert.assertEquals(buf.refCnt(), 1);
+            cbuf.release();
+            Assert.assertEquals(cbuf.refCnt(), 0);
+            Assert.assertEquals(buf.refCnt(), 0);
+        } finally {
+            ((DirectBuffer) buffer).cleaner().clean();
+        }
+    }
 }