fix(datastore): fix copyTo function of SwByteBuffer

star-whale · Sep 16, 2022 · b8d7c1c · b8d7c1c
1 parent ba0710e
commit b8d7c1c
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 1 deletion.
diff --git a/server/controller/src/main/java/ai/starwhale/mlops/memory/impl/SwByteBuffer.java b/server/controller/src/main/java/ai/starwhale/mlops/memory/impl/SwByteBuffer.java
@@ -122,6 +122,9 @@ public int getBytes(int index, byte[] b, int offset, int len) {
     @Override
     public void setBytes(int index, byte[] b, int offset, int len) {
         this.buf.position(index);
+        if (len > this.buf.remaining()) {
+            len = this.buf.remaining();
+        }
         this.buf.put(b, offset, len);
     }
 
@@ -141,7 +144,7 @@ public SwBuffer slice(int offset, int len) {
 
     @Override
     public void copyTo(SwBuffer buf) {
-        buf.setBytes(0, this.buf.array(), 0, this.buf.limit());
+        buf.setBytes(0, this.buf.array(), this.buf.arrayOffset(), this.buf.limit());
     }
 
     @Override

diff --git a/server/controller/src/test/java/ai/starwhale/mlops/memory/impl/SwByteBufferTest.java b/server/controller/src/test/java/ai/starwhale/mlops/memory/impl/SwByteBufferTest.java
@@ -19,6 +19,7 @@
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.is;
 
+import ai.starwhale.mlops.memory.SwBuffer;
 import java.nio.ByteBuffer;
 import java.nio.charset.StandardCharsets;
 import java.util.Arrays;
@@ -79,4 +80,31 @@ public void testSlice() {
         this.buffer.setString(0, "1234567890");
         assertThat(buf.asByteBuffer(), is(ByteBuffer.wrap("23456".getBytes(StandardCharsets.UTF_8))));
     }
+
+    @Test
+    public void testCopyTo() {
+        var base = new SwByteBuffer(5);
+        base.setBytes(0, new byte[]{1, 2, 3, 4, 5}, 0, 5);
+        // make a SwBuffer with non-zero offset
+        var src = base.slice(1, base.capacity() - 1);
+        var dst = new SwByteBuffer(5);
+        src.copyTo(dst);
+        assertThat(dst.getByte(0), is(src.getByte(0)));
+    }
+
+    @Test
+    void testSetBytes() {
+        final int cap = 3;
+        var buf = new SwByteBuffer(cap);
+        var src = new byte[]{1, 2, 3, 4};
+        buf.setBytes(0, src, 0, src.length);
+
+        var b = new byte[5];
+        assertThat(buf.getBytes(0, b, 0, b.length), is(cap));
+        assertThat(b, is(new byte[]{1, 2, 3, 0, 0}));
+
+        buf.setBytes(1, src, 0, src.length);
+        assertThat(buf.getBytes(0, b, 0, b.length), is(cap));
+        assertThat(b, is(new byte[]{1, 1, 2, 0, 0}));
+    }
 }