Merge pull request #289 from ipinfo/large-db-support

Add support for files larger than 2GB

Author: oschwald

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

@@ -0,0 +1,117 @@
+package com.maxmind.db;
+
+import java.io.IOException;
+import java.nio.channels.FileChannel;
+import java.nio.charset.CharacterCodingException;
+import java.nio.charset.CharsetDecoder;
+
+/**
+ * A generic buffer abstraction that supports sequential and random access
+ * to binary data. Implementations may be backed by a single {@link
+ * java.nio.ByteBuffer} or multiple buffers for larger capacities.
+ *
+ * <p>This interface is designed to provide a long-based API while
+ * remaining compatible with the limitations of underlying storage.
+ */
+interface Buffer {
+    /**
+     * Returns the total capacity of this buffer in bytes.
+     *
+     * @return the capacity
+     */
+    long capacity();
+
+    /**
+     * Returns the current position of this buffer.
+     *
+     * @return the position
+     */
+    long position();
+
+    /**
+     * Sets the buffer's position.
+     *
+     * @param newPosition the new position
+     * @return this buffer
+     */
+    Buffer position(long newPosition);
+
+    /**
+     * Returns the current limit of this buffer.
+     *
+     * @return the limit
+     */
+    long limit();
+
+    /**
+     * Sets the buffer's limit.
+     *
+     * @param newLimit the new limit
+     * @return this buffer
+     */
+    Buffer limit(long newLimit);
+
+    /**
+     * Reads the next byte at the current position and advances the position.
+     *
+     * @return the byte value
+     */
+    byte get();
+
+    /**
+     * Reads bytes into the given array and advances the position.
+     *
+     * @param dst the destination array
+     * @return this buffer
+     */
+    Buffer get(byte[] dst);
+
+    /**
+     * Reads a byte at the given absolute index without changing the position.
+     *
+     * @param index the index to read from
+     * @return the byte value
+     */
+    byte get(long index);
+
+    /**
+     * Reads the next 8 bytes as a double and advances the position.
+     *
+     * @return the double value
+     */
+    double getDouble();
+
+    /**
+     * Reads the next 4 bytes as a float and advances the position.
+     *
+     * @return the float value
+     */
+    float getFloat();
+
+    /**
+     * Creates a new buffer that shares the same content but has independent
+     * position, limit, and mark values.
+     *
+     * @return a duplicate buffer
+     */
+    Buffer duplicate();
+
+    /**
+     * Reads data from the given channel into this buffer starting at the
+     * current position.
+     *
+     * @param channel the file channel
+     * @return the number of bytes read
+     * @throws IOException if an I/O error occurs
+     */
+    long readFrom(FileChannel channel) throws IOException;
+
+    /**
+     * Decodes the buffer's content into a string using the given decoder.
+     *
+     * @param decoder the charset decoder
+     * @return the decoded string
+     * @throws CharacterCodingException if decoding fails
+     */
+    String decode(CharsetDecoder decoder) throws CharacterCodingException;
+}

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

@@ -1,34 +1,46 @@
 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.ByteBuffer;
 import java.nio.channels.FileChannel;
-import java.nio.channels.FileChannel.MapMode;
+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 ByteBuffer buffer;
+    private final Buffer buffer;
 
     BufferHolder(File database, FileMode mode) throws IOException {
-        try (
-            final RandomAccessFile file = new RandomAccessFile(database, "r");
-            final FileChannel channel = file.getChannel()
-        ) {
+        this(database, mode, MultiBuffer.DEFAULT_CHUNK_SIZE);
+    }
+
+    BufferHolder(File database, FileMode mode, int chunkSize) throws IOException {
+        try (RandomAccessFile file = new RandomAccessFile(database, "r");
+             FileChannel channel = file.getChannel()) {
+            long size = channel.size();
             if (mode == FileMode.MEMORY) {
-                final ByteBuffer buf = ByteBuffer.wrap(new byte[(int) channel.size()]);
-                if (channel.read(buf) != buf.capacity()) {
+                Buffer buf;
+                if (size <= chunkSize) {
+                    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.asReadOnlyBuffer();
+                this.buffer = buf;
             } else {
-                this.buffer = channel.map(MapMode.READ_ONLY, 0, channel.size()).asReadOnlyBuffer();
+                if (size <= chunkSize) {
+                    this.buffer = SingleBuffer.mapFromChannel(channel);
+                } else {
+                    this.buffer = MultiBuffer.mapFromChannel(channel);
+                }
             }
         }
     }
@@ -41,23 +53,44 @@ final class BufferHolder {
      * @throws NullPointerException if you provide a NULL InputStream
      */
     BufferHolder(InputStream stream) throws IOException {
+        this(stream, MultiBuffer.DEFAULT_CHUNK_SIZE);
+    }
+
+    BufferHolder(InputStream stream, int chunkSize) throws  IOException {
         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);
+        List<ByteBuffer> chunks = new ArrayList<>();
+        long total = 0;
+        byte[] tmp = new byte[chunkSize];
+        int read;
+
+        while (-1 != (read = stream.read(tmp))) {
+            ByteBuffer chunk = ByteBuffer.allocate(read);
+            chunk.put(tmp, 0, read);
+            chunk.flip();
+            chunks.add(chunk);
+            total += read;
+        }
+
+        if (total <= chunkSize) {
+            byte[] data = new byte[(int) total];
+            int pos = 0;
+            for (ByteBuffer chunk : chunks) {
+                System.arraycopy(chunk.array(), 0, data, pos, chunk.capacity());
+                pos += chunk.capacity();
+            }
+            this.buffer = SingleBuffer.wrap(data);
+        } else {
+            this.buffer = new MultiBuffer(chunks.toArray(new ByteBuffer[0]), chunkSize);
         }
-        this.buffer = ByteBuffer.wrap(baos.toByteArray()).asReadOnlyBuffer();
     }
 
     /*
-     * Returns a duplicate of the underlying ByteBuffer. The returned ByteBuffer
+     * Returns a duplicate of the underlying Buffer. The returned Buffer
      * should not be shared between threads.
      */
-    ByteBuffer get() {
+    Buffer get() {
         // The Java API docs for buffer state:
         //
         //     Buffers are not safe for use by multiple concurrent threads. If a buffer is to be
@@ -70,7 +103,7 @@ ByteBuffer get() {
         // * https://github.com/maxmind/MaxMind-DB-Reader-java/issues/65
         // * https://github.com/maxmind/MaxMind-DB-Reader-java/pull/69
         //
-        // Given that we are not modifying the original ByteBuffer in any way and all currently
+        // Given that we are not modifying the original Buffer in any way and all currently
         // known and most reasonably imaginable implementations of duplicate() only do read
         // operations on the original buffer object, the risk of not synchronizing this call seems
         // relatively low and worth taking for the performance benefit when lookups are being done

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

@@ -10,5 +10,5 @@
  * @param cls the class of the value
  * @param type the type of the value
  */
-public record CacheKey<T>(int offset, Class<T> cls, java.lang.reflect.Type type) {
+public record CacheKey<T>(long offset, Class<T> cls, java.lang.reflect.Type type) {
 }

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

@@ -1,4 +1,4 @@
 package com.maxmind.db;
 
-record CtrlData(Type type, int ctrlByte, int offset, int size) {
+record CtrlData(Type type, int ctrlByte, long offset, int size) {
 }

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

@@ -34,11 +34,11 @@ class Decoder {
 
     private final CharsetDecoder utfDecoder = UTF_8.newDecoder();
 
-    private final ByteBuffer buffer;
+    private final Buffer buffer;
 
     private final ConcurrentHashMap<Class<?>, CachedConstructor<?>> constructors;
 
-    Decoder(NodeCache cache, ByteBuffer buffer, long pointerBase) {
+    Decoder(NodeCache cache, Buffer buffer, long pointerBase) {
         this(
             cache,
             buffer,
@@ -49,7 +49,7 @@ class Decoder {
 
     Decoder(
         NodeCache cache,
-        ByteBuffer buffer,
+        Buffer buffer,
         long pointerBase,
         ConcurrentHashMap<Class<?>, CachedConstructor<?>> constructors
     ) {
@@ -61,7 +61,7 @@ class Decoder {
 
     private final NodeCache.Loader cacheLoader = this::decode;
 
-    <T> T decode(int offset, Class<T> cls) throws IOException {
+    <T> T decode(long offset, Class<T> cls) throws IOException {
         if (offset >= this.buffer.capacity()) {
             throw new InvalidDatabaseException(
                 "The MaxMind DB file's data section contains bad data: "
@@ -73,7 +73,7 @@ <T> T decode(int offset, Class<T> cls) throws IOException {
     }
 
     private <T> DecodedValue decode(CacheKey<T> key) throws IOException {
-        int offset = key.offset();
+        long offset = key.offset();
         if (offset >= this.buffer.capacity()) {
             throw new InvalidDatabaseException(
                 "The MaxMind DB file's data section contains bad data: "
@@ -132,8 +132,8 @@ private <T> DecodedValue decode(Class<T> cls, java.lang.reflect.Type genericType
 
     DecodedValue decodePointer(long pointer, Class<?> cls, java.lang.reflect.Type genericType)
             throws IOException {
-        int targetOffset = (int) pointer;
-        int position = buffer.position();
+        long targetOffset = pointer;
+        long position = buffer.position();
 
         CacheKey<?> key = new CacheKey<>(targetOffset, cls, genericType);
         DecodedValue o = cache.get(key, cacheLoader);
@@ -185,10 +185,10 @@ private <T> Object decodeByType(
         }
     }
 
-    private String decodeString(int size) throws CharacterCodingException {
-        int oldLimit = buffer.limit();
+    private String decodeString(long size) throws CharacterCodingException {
+        long oldLimit = buffer.limit();
         buffer.limit(buffer.position() + size);
-        String s = utfDecoder.decode(buffer).toString();
+        String s = buffer.decode(utfDecoder);
         buffer.limit(oldLimit);
         return s;
     }
@@ -202,9 +202,13 @@ private int decodeInt32(int size) {
     }
 
     private long decodeLong(int size) {
-        long integer = 0;
+        return Decoder.decodeLong(this.buffer, 0, size);
+    }
+
+    static long decodeLong(Buffer buffer, int base, int size) {
+        long integer = base;
         for (int i = 0; i < size; i++) {
-            integer = (integer << 8) | (this.buffer.get() & 0xFF);
+            integer = (integer << 8) | (buffer.get() & 0xFF);
         }
         return integer;
     }
@@ -221,7 +225,7 @@ private int decodeInteger(int base, int size) {
         return Decoder.decodeInteger(this.buffer, base, size);
     }
 
-    static int decodeInteger(ByteBuffer buffer, int base, int size) {
+    static int decodeInteger(Buffer buffer, int base, int size) {
         int integer = base;
         for (int i = 0; i < size; i++) {
             integer = (integer << 8) | (buffer.get() & 0xFF);
@@ -412,7 +416,7 @@ private <T> Object decodeMapIntoObject(int size, Class<T> cls)
 
             Integer parameterIndex = parameterIndexes.get(key);
             if (parameterIndex == null) {
-                int offset = this.nextValueOffset(this.buffer.position(), 1);
+                long offset = this.nextValueOffset(this.buffer.position(), 1);
                 this.buffer.position(offset);
                 continue;
             }
@@ -485,7 +489,7 @@ private static <T> String getParameterName(
                 + " is not annotated with MaxMindDbParameter.");
     }
 
-    private int nextValueOffset(int offset, int numberToSkip)
+    private long nextValueOffset(long offset, int numberToSkip)
         throws InvalidDatabaseException {
         if (numberToSkip == 0) {
             return offset;
@@ -518,7 +522,7 @@ private int nextValueOffset(int offset, int numberToSkip)
         return nextValueOffset(offset, numberToSkip - 1);
     }
 
-    private CtrlData getCtrlData(int offset)
+    private CtrlData getCtrlData(long offset)
         throws InvalidDatabaseException {
         if (offset >= this.buffer.capacity()) {
             throw new InvalidDatabaseException(
@@ -566,7 +570,7 @@ private byte[] getByteArray(int length) {
         return Decoder.getByteArray(this.buffer, length);
     }
 
-    private static byte[] getByteArray(ByteBuffer buffer, int length) {
+    private static byte[] getByteArray(Buffer buffer, int length) {
         byte[] bytes = new byte[length];
         buffer.get(bytes);
         return bytes;