Use MultiBuffer when DB size can't fit in SingleBuffer

Author: silvanocerza

src/main/java/com/maxmind/db/BufferHolder.java

@@ -1,32 +1,41 @@
 package com.maxmind.db;
 
 import com.maxmind.db.Reader.FileMode;
-import java.io.ByteArrayOutputStream;
 import java.io.File;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.RandomAccessFile;
 import java.nio.channels.FileChannel;
+import java.util.ArrayList;
+import java.util.List;
 
 final class BufferHolder {
     // DO NOT PASS OUTSIDE THIS CLASS. Doing so will remove thread safety.
     private final Buffer buffer;
 
     BufferHolder(File database, FileMode mode) throws IOException {
-        try (
-            final RandomAccessFile file = new RandomAccessFile(database, "r");
-            final FileChannel channel = file.getChannel()
-        ) {
+        try (RandomAccessFile file = new RandomAccessFile(database, "r");
+             FileChannel channel = file.getChannel()) {
+            long size = channel.size();
             if (mode == FileMode.MEMORY) {
-                final SingleBuffer buf = SingleBuffer.wrap(new byte[(int) channel.size()]);
+                Buffer buf;
+                if (size <= Integer.MAX_VALUE) {
+                    buf = new SingleBuffer(size);
+                } else {
+                    buf = new MultiBuffer(size);
+                }
                 if (buf.readFrom(channel) != buf.capacity()) {
                     throw new IOException("Unable to read "
-                        + database.getName()
-                        + " into memory. Unexpected end of stream.");
+                            + database.getName()
+                            + " into memory. Unexpected end of stream.");
                 }
                 this.buffer = buf;
             } else {
-                this.buffer = SingleBuffer.mapFromChannel(channel);
+                if (size <= Integer.MAX_VALUE) {
+                    this.buffer = SingleBuffer.mapFromChannel(channel);
+                } else {
+                    this.buffer = MultiBuffer.mapFromChannel(channel);
+                }
             }
         }
     }
@@ -42,13 +51,30 @@ final class BufferHolder {
         if (null == stream) {
             throw new NullPointerException("Unable to use a NULL InputStream");
         }
-        final ByteArrayOutputStream baos = new ByteArrayOutputStream();
-        final byte[] bytes = new byte[16 * 1024];
-        int br;
-        while (-1 != (br = stream.read(bytes))) {
-            baos.write(bytes, 0, br);
+        final int CHUNK_SIZE = 16 * 1024;
+        List<byte[]> chunks = new ArrayList<>();
+        long total = 0;
+        int read;
+        byte[] tmp = new byte[CHUNK_SIZE];
+
+        while (-1 != (read = stream.read(tmp))) {
+            byte[] copy = new byte[read];
+            System.arraycopy(tmp, 0, copy, 0, read);
+            chunks.add(copy);
+            total += read;
+        }
+
+        if (total <= Integer.MAX_VALUE) {
+            byte[] data = new byte[(int) total];
+            int pos = 0;
+            for (byte[] chunk : chunks) {
+                System.arraycopy(chunk, 0, data, pos, chunk.length);
+                pos += chunk.length;
+            }
+            this.buffer = SingleBuffer.wrap(data);
+        } else {
+            this.buffer = MultiBuffer.wrap(chunks);
         }
-        this.buffer = SingleBuffer.wrap(baos.toByteArray());
     }
 
     /*

src/main/java/com/maxmind/db/Metadata.java

@@ -24,11 +24,11 @@ public final class Metadata {
 
     private final int nodeByteSize;
 
-    private final int nodeCount;
+    private final long nodeCount;
 
     private final int recordSize;
 
-    private final int searchTreeSize;
+    private final long searchTreeSize;
 
     /**
      * Constructs a {@code Metadata} object.
@@ -71,7 +71,7 @@ public Metadata(
         this.languages = languages;
         this.description = description;
         this.ipVersion = ipVersion;
-        this.nodeCount = (int) nodeCount;
+        this.nodeCount = nodeCount;
         this.recordSize = recordSize;
 
         this.nodeByteSize = this.recordSize / 4;
@@ -140,7 +140,7 @@ int getNodeByteSize() {
     /**
      * @return the number of nodes in the search tree.
      */
-    int getNodeCount() {
+    long getNodeCount() {
         return this.nodeCount;
     }
 
@@ -155,7 +155,7 @@ int getRecordSize() {
     /**
      * @return the searchTreeSize
      */
-    int getSearchTreeSize() {
+    long getSearchTreeSize() {
         return this.searchTreeSize;
     }
 

src/main/java/com/maxmind/db/Networks.java

@@ -173,7 +173,7 @@ static class NetworkNode {
         /** The prefix of the node. */
         public int prefix;
         /** The node number. */
-        public int pointer;
+        public long pointer;
 
         /**
          * Constructs a network node for internal use.
@@ -182,7 +182,7 @@ static class NetworkNode {
          * @param prefix The prefix of the node.
          * @param pointer The node number
          */
-        NetworkNode(byte[] ip, int prefix, int pointer) {
+        NetworkNode(byte[] ip, int prefix, long pointer) {
             this.ip = ip;
             this.prefix = prefix;
             this.pointer = pointer;

src/main/java/com/maxmind/db/Reader.java

@@ -22,7 +22,7 @@ public final class Reader implements Closeable {
         (byte) 0xCD, (byte) 0xEF, 'M', 'a', 'x', 'M', 'i', 'n', 'd', '.',
         'c', 'o', 'm'};
 
-    private final int ipV4Start;
+    private final long ipV4Start;
     private final Metadata metadata;
     private final AtomicReference<BufferHolder> bufferHolderReference;
     private final NodeCache cache;
@@ -152,7 +152,7 @@ public <T> T get(InetAddress ipAddress, Class<T> cls) throws IOException {
         return getRecord(ipAddress, cls).getData();
     }
 
-    int getIpv4Start() {
+    long getIpv4Start() {
         return this.ipV4Start;
     }
 
@@ -171,12 +171,12 @@ public <T> DatabaseRecord<T> getRecord(InetAddress ipAddress, Class<T> cls)
 
         byte[] rawAddress = ipAddress.getAddress();
 
-        int[] traverseResult = traverseTree(rawAddress, rawAddress.length * 8);
+        long[] traverseResult = traverseTree(rawAddress, rawAddress.length * 8);
 
-        int pl = traverseResult[1];
-        int record = traverseResult[0];
+        long record = traverseResult[0];
+        int pl = (int) traverseResult[1];
 
-        int nodeCount = this.metadata.getNodeCount();
+        long nodeCount = this.metadata.getNodeCount();
         Buffer buffer = this.getBufferHolder().get();
         T dataRecord = null;
         if (record > nodeCount) {
@@ -253,7 +253,7 @@ BufferHolder getBufferHolder() throws ClosedDatabaseException {
         return bufferHolder;
     }
 
-    private int startNode(int bitLength) {
+    private long startNode(int bitLength) {
         // Check if we are looking up an IPv4 address in an IPv6 tree. If this
         // is the case, we can skip over the first 96 nodes.
         if (this.metadata.getIpVersion() == 6 && bitLength == 32) {
@@ -264,13 +264,13 @@ private int startNode(int bitLength) {
         return 0;
     }
 
-    private int findIpV4StartNode(Buffer buffer)
+    private long findIpV4StartNode(Buffer buffer)
         throws InvalidDatabaseException {
         if (this.metadata.getIpVersion() == 4) {
             return 0;
         }
 
-        int node = 0;
+        long node = 0;
         for (int i = 0; i < 96 && node < this.metadata.getNodeCount(); i++) {
             node = this.readNode(buffer, node, 0);
         }
@@ -319,9 +319,9 @@ public <T> Networks<T> networksWithin(
             prefixLength += 96;
         }
 
-        int[] traverseResult = this.traverseTree(ipBytes, prefixLength);
-        int node = traverseResult[0];
-        int prefix = traverseResult[1];
+        long[] traverseResult = this.traverseTree(ipBytes, prefixLength);
+        long node = traverseResult[0];
+        int prefix = (int) traverseResult[1];
 
         return new Networks<>(this, includeAliasedNetworks,
             new Networks.NetworkNode[] {new Networks.NetworkNode(ipBytes, prefix, node)},
@@ -335,12 +335,12 @@ public <T> Networks<T> networksWithin(
      * @param bitCount The prefix.
      * @return int[]
      */
-    private int[] traverseTree(byte[] ip, int bitCount)
+    private long[] traverseTree(byte[] ip, int bitCount)
         throws ClosedDatabaseException, InvalidDatabaseException {
         Buffer buffer = this.getBufferHolder().get();
         int bitLength = ip.length * 8;
-        int record = this.startNode(bitLength);
-        int nodeCount = this.metadata.getNodeCount();
+        long record = this.startNode(bitLength);
+        long nodeCount = this.metadata.getNodeCount();
 
         int i = 0;
         for (; i < bitCount && record < nodeCount; i++) {
@@ -352,19 +352,19 @@ int record = this.startNode(bitLength);
             record = this.readNode(buffer, record, bit);
         }
 
-        return new int[]{record, i};
+        return new long[]{record, i};
     }
 
-    int readNode(Buffer buffer, int nodeNumber, int index)
+    int readNode(Buffer buffer, long nodeNumber, int index)
             throws InvalidDatabaseException {
         // index is the index of the record within the node, which
         // can either be 0 or 1.
-        int baseOffset = nodeNumber * this.metadata.getNodeByteSize();
+        long baseOffset = nodeNumber * this.metadata.getNodeByteSize();
 
         switch (this.metadata.getRecordSize()) {
             case 24:
                 // For a 24 bit record, each record is 3 bytes.
-                buffer.position(baseOffset + index * 3);
+                buffer.position(baseOffset + (long) index * 3);
                 return Decoder.decodeInteger(buffer, 0, 3);
             case 28:
                 int middle = buffer.get(baseOffset + 3);
@@ -377,10 +377,10 @@ int readNode(Buffer buffer, int nodeNumber, int index)
                     // We get the most significant byte of the second record.
                     middle = 0x0F & middle;
                 }
-                buffer.position(baseOffset + index * 4);
+                buffer.position(baseOffset + (long) index * 4);
                 return Decoder.decodeInteger(buffer, middle, 3);
             case 32:
-                buffer.position(baseOffset + index * 4);
+                buffer.position(baseOffset + (long) index * 4);
                 return Decoder.decodeInteger(buffer, 0, 4);
             default:
                 throw new InvalidDatabaseException("Unknown record size: "
@@ -390,10 +390,10 @@ int readNode(Buffer buffer, int nodeNumber, int index)
 
     <T> T resolveDataPointer(
         Buffer buffer,
-        int pointer,
+        long pointer,
         Class<T> cls
     ) throws IOException {
-        int resolved = (pointer - this.metadata.getNodeCount())
+        long resolved = (pointer - this.metadata.getNodeCount())
             + this.metadata.getSearchTreeSize();
 
         if (resolved >= buffer.capacity()) {

src/test/java/com/maxmind/db/DecoderTest.java

@@ -8,7 +8,6 @@
 
 import java.io.IOException;
 import java.math.BigInteger;
-import java.nio.ByteBuffer;
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.HashMap;

src/test/java/com/maxmind/db/TestDecoder.java

@@ -5,7 +5,7 @@
 
 final class TestDecoder extends Decoder {
 
-    TestDecoder(NodeCache cache, SingleBuffer buffer, long pointerBase) {
+    TestDecoder(NodeCache cache, Buffer buffer, long pointerBase) {
         super(cache, buffer, pointerBase);
     }