JAVA-5789 Implements batch ByteBuf::indexOf

franz1981 · franz1981 · commit 816ddcddf638 · 2025-02-12T14:22:10.000+01:00
diff --git a/bson/src/main/org/bson/ByteBuf.java b/bson/src/main/org/bson/ByteBuf.java
@@ -168,6 +168,23 @@ public interface ByteBuf  {
      */
     ByteBuf clear();
 
+
+    /**
+     * Returns the index within this buffer of the first occurrence of the given byte
+     *
+     * @param b     The byte to search for
+     * @return The index within this buffer of the first occurrence of the given byte, starting the search at the specified index, or {@code
+     * -1} if the byte does not occur before the limit
+     */
+    default int indexOf(byte b) {
+        for (int i = position(); i < limit(); i++) {
+            if (get(i) == b) {
+                return i;
+            }
+        }
+        return -1;
+    }
+
     /**
      * Modifies this buffer's byte order.
      *
diff --git a/bson/src/main/org/bson/ByteBufNIO.java b/bson/src/main/org/bson/ByteBufNIO.java
@@ -64,6 +64,31 @@ public void release() {
         }
     }
 
+    @Override
+    public int indexOf(final byte b) {
+        ByteBuffer buf = this.buf;
+        // readonly buffers won't go into the fast-path
+        if (buf.hasArray()) {
+            byte[] array = buf.array();
+            int offset = buf.arrayOffset() + buf.position();
+            int length = buf.remaining();
+            for (int i = 0; i < length; i++) {
+                if (array[offset + i] == b) {
+                    return i;
+                }
+            }
+            return -1;
+        }
+        int position = buf.position();
+        int limit = buf.limit();
+        for (int i = position; i < limit; i++) {
+            if (buf.get(i) == b) {
+                return i - position;
+            }
+        }
+        return -1;
+    }
+
     @Override
     public int capacity() {
         return buf.capacity();
diff --git a/bson/src/main/org/bson/io/ByteBufferBsonInput.java b/bson/src/main/org/bson/io/ByteBufferBsonInput.java
@@ -159,13 +159,12 @@ private String readString(final int size) {
     @Override
     public void skipCString() {
         ensureOpen();
-        boolean checkNext = true;
-        while (checkNext) {
-            if (!buffer.hasRemaining()) {
-                throw new BsonSerializationException("Found a BSON string that is not null-terminated");
-            }
-            checkNext = buffer.get() != 0;
+        int indexOfZero = buffer.indexOf((byte) 0);
+        if (indexOfZero == -1) {
+            buffer.position(buffer.limit());
+            throw new BsonSerializationException("Found a BSON string that is not null-terminated");
         }
+        buffer.position(indexOfZero + 1);
     }
 
     @Override
diff --git a/driver-core/src/main/com/mongodb/internal/connection/CompositeByteBuf.java b/driver-core/src/main/com/mongodb/internal/connection/CompositeByteBuf.java
@@ -42,7 +42,9 @@ class CompositeByteBuf implements ByteBuf {
 
         int offset = 0;
         for (ByteBuf cur : buffers) {
-            Component component = new Component(cur.asReadOnly().order(ByteOrder.LITTLE_ENDIAN), offset);
+            // since we don't expose any method to modify the buffer nor expose its array we can avoid using read-only ones
+            // to speed up ByteBufNIO::indexOf, is it's the concrete type used
+            Component component = new Component(cur.duplicate().order(ByteOrder.LITTLE_ENDIAN), offset);
             components.add(component);
             offset = component.endOffset;
         }
@@ -63,6 +65,18 @@ public ByteBuf order(final ByteOrder byteOrder) {
         return this;
     }
 
+    @Override
+    public int indexOf(final byte b) {
+        // use this pattern to save creating a new iterator
+        for (int i = 0, size = components.size(); i < size; i++) {
+            int index = components.get(i).indexOf(b);
+            if (index != -1) {
+                return index;
+            }
+        }
+        return -1;
+    }
+
     @Override
     public int capacity() {
         return components.get(components.size() - 1).endOffset;
@@ -340,5 +354,13 @@ private static final class Component {
             this.offset = offset;
             this.endOffset = offset + length;
         }
+
+        public int indexOf(final byte b) {
+            int i = buffer.indexOf(b);
+            if (i != -1) {
+                return i + offset;
+            }
+            return -1;
+        }
     }
 }
diff --git a/driver-core/src/main/com/mongodb/internal/connection/netty/NettyByteBuf.java b/driver-core/src/main/com/mongodb/internal/connection/netty/NettyByteBuf.java
@@ -89,6 +89,16 @@ public ByteBuf put(final byte b) {
         return this;
     }
 
+    @Override
+    public int indexOf(final byte b) {
+        int position = isWriting ? proxied.writerIndex() : proxied.readerIndex();
+        int limit = position + (isWriting ? proxied.writableBytes() : proxied.readableBytes());
+        if (position == limit) {
+            return -1;
+        }
+        return proxied.indexOf(position, limit, b);
+    }
+
     @Override
     public ByteBuf flip() {
         isWriting = !isWriting;
diff --git a/driver-core/src/test/unit/com/mongodb/internal/connection/ByteBufSpecification.groovy b/driver-core/src/test/unit/com/mongodb/internal/connection/ByteBufSpecification.groovy
@@ -24,6 +24,56 @@ import org.bson.ByteBuf
 import spock.lang.Specification
 
 class ByteBufSpecification extends Specification {
+
+
+    def "should find the first occurrence of a byte"() {
+        given:
+        def buffer = provider.getBuffer(1024)
+
+        when:
+        buffer.with {
+            put((byte) 1)
+            put((byte) 2)
+            put((byte) 3)
+            put((byte) 4)
+            put((byte) 5)
+            flip()
+        }
+
+        then:
+        buffer.indexOf((byte) 3) == 2
+
+        cleanup:
+        buffer.release()
+
+        where:
+        provider << [new NettyBufferProvider(), new SimpleBufferProvider()]
+    }
+
+    def "should not find any occurrence of a byte"() {
+            given:
+        def buffer = provider.getBuffer(1024)
+
+        when:
+        buffer.with {
+            put((byte) 1)
+            put((byte) 2)
+            put((byte) 3)
+            put((byte) 4)
+            put((byte) 5)
+            flip()
+        }
+
+        then:
+        buffer.indexOf((byte) 6) == -1
+
+        cleanup:
+        buffer.release()
+
+        where:
+        provider << [new NettyBufferProvider(), new SimpleBufferProvider()]
+    }
+
     def 'should put a byte'() {
         given:
         def buffer = provider.getBuffer(1024)
diff --git a/driver-core/src/test/unit/com/mongodb/internal/connection/CompositeByteBufSpecification.groovy b/driver-core/src/test/unit/com/mongodb/internal/connection/CompositeByteBufSpecification.groovy
@@ -290,6 +290,37 @@ class CompositeByteBufSpecification extends Specification {
         !buf.hasRemaining()
     }
 
+    def "should find the first occurrence of a byte"() {
+        given:
+        def buf = new CompositeByteBuf([
+                        new ByteBufNIO(ByteBuffer.wrap([5, 5, 5, 5] as byte[])),
+                        new ByteBufNIO(ByteBuffer.wrap([1, 2, 3, 4] as byte[]))])
+        when:
+        byte b = 3;
+
+        then:
+        buf.indexOf(b) == 6
+
+        cleanup:
+        buf.release()
+    }
+
+    def "should not find any occurrence of a byte"() {
+        given:
+        def buf = new CompositeByteBuf([
+                new ByteBufNIO(ByteBuffer.wrap([5, 5, 5, 5] as byte[])),
+                new ByteBufNIO(ByteBuffer.wrap([1, 2, 3, 4] as byte[]))])
+
+        when:
+          byte b = 6;
+
+        then:
+        buf.indexOf(b) == -1
+
+        cleanup:
+        buf.release()
+    }
+
     def 'absolute getInt should read little endian integer and preserve position'() {
         given:
         def byteBuffer = new ByteBufNIO(ByteBuffer.wrap([1, 2, 3, 4] as byte[]))
