Merge pull request #368 from maxmind/greg/stf-383

STF-383: Fix off-heap memory growth in FileMode.MEMORY

Author: horgh

CHANGELOG.md

@@ -1,6 +1,25 @@
 CHANGELOG
 =========
 
+4.1.0
+------------------
+
+* Fixed unbounded off-heap memory growth when initializing the reader in
+  `FileMode.MEMORY`. The previous implementation read the database via
+  `FileChannel.read()` into a heap buffer, which causes the JDK to cache
+  temporary direct ByteBuffers in per-thread storage
+  (`sun.nio.ch.Util.BufferCache`). Repeated initialization across different
+  threads could grow this cache without bound. The reader now uses
+  `FileInputStream` for `MEMORY` mode, which bypasses the cache.
+  `FileMode.MEMORY_MAPPED` was unaffected.
+* Fixed a latent short-read bug in the multi-chunk `FileMode.MEMORY`
+  load path introduced in 4.0.0. `FileChannel.read(ByteBuffer)` is not
+  contractually obligated to fully fill the destination buffer; a
+  short read could have caused silent truncation of an in-memory
+  database. Affects databases larger than ~2GB (the default chunk
+  size). The new chunked read loop retries until each chunk is fully
+  populated.
+
 4.0.2 (2025-12-08)
 ------------------
 

mise.lock

@@ -0,0 +1,18 @@
+[settings]
+experimental = true
+lockfile = true
+disable_backends = [
+    "asdf",
+    "vfox",
+]
+
+[tools]
+java = "latest"
+maven = "latest"
+
+[hooks]
+enter = "mise install --quiet --locked"
+
+[[watch_files]]
+patterns = ["mise.toml", "mise.lock"]
+run = "mise install --quiet --locked"

mise.toml

@@ -3,6 +3,7 @@
 import com.maxmind.db.Reader.FileMode;
 import java.io.ByteArrayOutputStream;
 import java.io.File;
+import java.io.FileInputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.RandomAccessFile;
@@ -23,53 +24,38 @@ final class BufferHolder {
     }
 
     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) {
+        if (mode == FileMode.MEMORY) {
+            // FileInputStream avoids the per-thread direct ByteBuffer cache that
+            // FileChannel.read() populates when reading into a heap buffer. That cache
+            // retains the largest direct buffer ever requested — under chunked MEMORY
+            // mode that would mean chunkSize bytes of off-heap memory held per loader
+            // thread for the JVM's lifetime.
+            try (FileInputStream stream = new FileInputStream(database)) {
+                // Size from the open fd (fstat) so it's atomic with the open;
+                // getChannel().size() does not populate the buffer cache.
+                long size = stream.getChannel().size();
+                var name = database.getName();
                 if (size <= chunkSize) {
-                    // Allocate, read, and make read-only
-                    ByteBuffer buffer = ByteBuffer.allocate((int) size);
-                    if (channel.read(buffer) != size) {
-                        throw new IOException("Unable to read "
-                                + database.getName()
-                                + " into memory. Unexpected end of stream.");
-                    }
-                    buffer.flip();
-                    this.buffer = new SingleBuffer(buffer);
+                    this.buffer = SingleBuffer.wrap(readFully(stream, (int) size, name));
                 } else {
-                    // Allocate chunks, read, and make read-only
                     var fullChunks = (int) (size / chunkSize);
                     var remainder = (int) (size % chunkSize);
                     var totalChunks = fullChunks + (remainder > 0 ? 1 : 0);
                     var buffers = new ByteBuffer[totalChunks];
-
                     for (int i = 0; i < fullChunks; i++) {
-                        buffers[i] = ByteBuffer.allocate(chunkSize);
+                        buffers[i] = ByteBuffer.wrap(readFully(stream, chunkSize, name));
                     }
                     if (remainder > 0) {
-                        buffers[totalChunks - 1] = ByteBuffer.allocate(remainder);
+                        buffers[totalChunks - 1] = ByteBuffer.wrap(
+                                readFully(stream, remainder, name));
                     }
-
-                    var totalRead = 0L;
-                    for (var buffer : buffers) {
-                        var read = channel.read(buffer);
-                        if (read == -1) {
-                            break;
-                        }
-                        totalRead += read;
-                        buffer.flip();
-                    }
-
-                    if (totalRead != size) {
-                        throw new IOException("Unable to read "
-                                + database.getName()
-                                + " into memory. Unexpected end of stream.");
-                    }
-
                     this.buffer = new MultiBuffer(buffers, chunkSize);
                 }
-            } else {
+            }
+        } else {
+            try (RandomAccessFile file = new RandomAccessFile(database, "r");
+                 FileChannel channel = file.getChannel()) {
+                long size = channel.size();
                 if (size <= chunkSize) {
                     this.buffer = SingleBuffer.mapFromChannel(channel);
                 } else {
@@ -79,11 +65,32 @@ final class BufferHolder {
         }
     }
 
-    BufferHolder(InputStream stream, int chunkSize) throws  IOException {
+    BufferHolder(InputStream stream, int chunkSize) throws IOException {
         if (null == stream) {
             throw new NullPointerException("Unable to use a NULL InputStream");
         }
+        this.buffer = readFromStream(stream, chunkSize);
+    }
+
+    // Pre-allocates exactly len bytes. Used by file-backed MEMORY mode where the size is
+    // known up front, avoiding the transient peak from ByteArrayOutputStream.grow() and
+    // the defensive copy in toByteArray().
+    private static byte[] readFully(InputStream stream, int len, String name) throws IOException {
+        var data = new byte[len];
+        var totalRead = 0;
+        while (totalRead < len) {
+            var n = stream.read(data, totalRead, len - totalRead);
+            if (n < 0) {
+                throw new IOException("Unable to read "
+                        + name
+                        + " into memory. Unexpected end of stream.");
+            }
+            totalRead += n;
+        }
+        return data;
+    }
 
+    private static Buffer readFromStream(InputStream stream, int chunkSize) throws IOException {
         // Read data from the stream in chunks to support databases >2GB.
         // Invariant: All chunks except the last are exactly chunkSize bytes.
         var chunks = new ArrayList<byte[]>();
@@ -116,17 +123,16 @@ final class BufferHolder {
 
         if (chunks.size() == 1) {
             // For databases that fit in a single chunk, use SingleBuffer
-            this.buffer = SingleBuffer.wrap(chunks.get(0));
-        } else {
-            // For large databases, wrap chunks in ByteBuffers and use MultiBuffer
-            // Guaranteed: chunks[0..n-2] all have length == chunkSize
-            // chunks[n-1] may have length < chunkSize
-            var buffers = new ByteBuffer[chunks.size()];
-            for (var i = 0; i < chunks.size(); i++) {
-                buffers[i] = ByteBuffer.wrap(chunks.get(i));
-            }
-            this.buffer = new MultiBuffer(buffers, chunkSize);
+            return SingleBuffer.wrap(chunks.get(0));
+        }
+        // For large databases, wrap chunks in ByteBuffers and use MultiBuffer
+        // Guaranteed: chunks[0..n-2] all have length == chunkSize
+        // chunks[n-1] may have length < chunkSize
+        var buffers = new ByteBuffer[chunks.size()];
+        for (var i = 0; i < chunks.size(); i++) {
+            buffers[i] = ByteBuffer.wrap(chunks.get(i));
         }
+        return new MultiBuffer(buffers, chunkSize);
     }
 
     /*

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

@@ -59,8 +59,12 @@ public Reader(File database) throws IOException {
     }
 
     Reader(File database, int chunkSize) throws IOException {
+        this(database, FileMode.MEMORY_MAPPED, chunkSize);
+    }
+
+    Reader(File database, FileMode fileMode, int chunkSize) throws IOException {
         this(
-            new BufferHolder(database, FileMode.MEMORY_MAPPED, chunkSize),
+            new BufferHolder(database, fileMode, chunkSize),
             database.getName(),
             NoCache.getInstance()
         );

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

@@ -12,6 +12,7 @@
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
+import com.maxmind.db.Reader.FileMode;
 import java.io.File;
 import java.io.IOException;
 import java.io.InputStream;
@@ -83,6 +84,24 @@ public void test(int chunkSize) throws IOException {
         }
     }
 
+    @ParameterizedTest
+    @MethodSource("chunkSizes")
+    public void testMemoryMode(int chunkSize) throws IOException {
+        for (long recordSize : new long[] {24, 28, 32}) {
+            for (int ipVersion : new int[] {4, 6}) {
+                var file = getFile("MaxMind-DB-test-ipv" + ipVersion + "-" + recordSize + ".mmdb");
+                try (var reader = new Reader(file, FileMode.MEMORY, chunkSize)) {
+                    this.testMetadata(reader, ipVersion, recordSize);
+                    if (ipVersion == 4) {
+                        this.testIpV4(reader, file);
+                    } else {
+                        this.testIpV6(reader, file);
+                    }
+                }
+            }
+        }
+    }
+
     static class GetRecordTest {
         InetAddress ip;
         File db;