Fix resource leak when using a thread pool

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

@@ -12,14 +12,14 @@
 
 import com.maxmind.db.Reader.FileMode;
 
-final class ThreadBuffer extends ThreadLocal<ByteBuffer> implements Closeable {
+final class BufferHolder implements Closeable {
     // DO NOT PASS THESE OUTSIDE THIS CLASS. Doing so will remove thread
     // safety.
     private final ByteBuffer buffer;
     private final RandomAccessFile raf;
     private final FileChannel fc;
 
-    ThreadBuffer(File database, FileMode mode) throws IOException {
+    BufferHolder(File database, FileMode mode) throws IOException {
         this.raf = new RandomAccessFile(database, "r");
         this.fc = this.raf.getChannel();
         if (mode == FileMode.MEMORY) {
@@ -41,7 +41,7 @@ final class ThreadBuffer extends ThreadLocal<ByteBuffer> implements Closeable {
      * @throws NullPointerException
      *             if you provide a NULL InputStream
      */
-    ThreadBuffer(InputStream stream) throws IOException {
+    BufferHolder(InputStream stream) throws IOException {
         if (null == stream) {
             throw new NullPointerException("Unable to use a NULL InputStream");
         }
@@ -57,14 +57,17 @@ final class ThreadBuffer extends ThreadLocal<ByteBuffer> implements Closeable {
     }
 
     // This is just to ease unit testing
-    ThreadBuffer(ByteBuffer buffer) {
+    BufferHolder(ByteBuffer buffer) {
         this.buffer = buffer;
         this.raf = null;
         this.fc = null;
     }
 
-    @Override
-    protected synchronized ByteBuffer initialValue() {
+    /*
+     * Returns a duplicate of the underlying ByteBuffer. The returned ByteBuffer
+     * should not be shared between threads.
+     */
+    synchronized ByteBuffer get() {
         return this.buffer.duplicate();
     }
 

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

@@ -18,6 +18,11 @@
 import com.fasterxml.jackson.databind.node.ObjectNode;
 import com.fasterxml.jackson.databind.node.TextNode;
 
+/*
+ * Decoder for MaxMind DB data.
+ *
+ * This class CANNOT be shared between threads
+ */
 final class Decoder {
     // XXX - This is only for unit testings. We should possibly make a
     // constructor to set this
@@ -26,7 +31,7 @@ final class Decoder {
 
     private final ObjectMapper objectMapper;
 
-    private final ThreadBuffer threadBuffer;
+    private final ByteBuffer buffer;
 
     static enum Type {
         EXTENDED, POINTER, UTF8_STRING, DOUBLE, BYTES, UINT16, UINT32, MAP, INT32, UINT64, UINT128, ARRAY, CONTAINER, END_MARKER, BOOLEAN, FLOAT;
@@ -73,22 +78,21 @@ void setOffset(int offset) {
 
     }
 
-    Decoder(ThreadBuffer threadBuffer, long pointerBase) {
+    Decoder(ByteBuffer buffer, long pointerBase) {
         this.pointerBase = pointerBase;
-        this.threadBuffer = threadBuffer;
+        this.buffer = buffer;
         this.objectMapper = new ObjectMapper();
     }
 
     Result decode(int offset) throws IOException {
-        ByteBuffer buffer = this.threadBuffer.get();
-        if (offset >= buffer.capacity()) {
+        if (offset >= this.buffer.capacity()) {
             throw new InvalidDatabaseException(
                     "The MaxMind DB file's data section contains bad data: "
                             + "pointer larger than the database.");
         }
 
-        buffer.position(offset);
-        int ctrlByte = 0xFF & buffer.get();
+        this.buffer.position(offset);
+        int ctrlByte = 0xFF & this.buffer.get();
         offset++;
 
         Type type = Type.fromControlByte(ctrlByte);
@@ -109,7 +113,7 @@ Result decode(int offset) throws IOException {
         }
 
         if (type.equals(Type.EXTENDED)) {
-            int nextByte = buffer.get();
+            int nextByte = this.buffer.get();
 
             int typeNum = nextByte + 7;
 
@@ -188,7 +192,7 @@ private Result decodePointer(int ctrlByte, int offset) {
     }
 
     private String decodeString(int size) {
-        ByteBuffer buffer = this.threadBuffer.get().slice();
+        ByteBuffer buffer = this.buffer.slice();
         buffer.limit(size);
         return Charset.forName("UTF-8").decode(buffer).toString();
     }
@@ -202,10 +206,9 @@ private IntNode decodeInt32(int size) {
     }
 
     private long decodeLong(int size) {
-        ByteBuffer buffer = this.threadBuffer.get();
         long integer = 0;
         for (int i = 0; i < size; i++) {
-            integer = (integer << 8) | (buffer.get() & 0xFF);
+            integer = (integer << 8) | (this.buffer.get() & 0xFF);
         }
         return integer;
     }
@@ -219,8 +222,7 @@ private int decodeInteger(int size) {
     }
 
     private int decodeInteger(int base, int size) {
-        ByteBuffer buffer = this.threadBuffer.get();
-        return Decoder.decodeInteger(buffer, base, size);
+        return Decoder.decodeInteger(this.buffer, base, size);
     }
 
     static int decodeInteger(ByteBuffer buffer, int base, int size) {
@@ -242,7 +244,7 @@ private DoubleNode decodeDouble(int size) throws InvalidDatabaseException {
                     "The MaxMind DB file's data section contains bad data: "
                             + "invalid size of double.");
         }
-        return new DoubleNode(this.threadBuffer.get().getDouble());
+        return new DoubleNode(this.buffer.getDouble());
     }
 
     private FloatNode decodeFloat(int size) throws InvalidDatabaseException {
@@ -251,7 +253,7 @@ private FloatNode decodeFloat(int size) throws InvalidDatabaseException {
                     "The MaxMind DB file's data section contains bad data: "
                             + "invalid size of float.");
         }
-        return new FloatNode(this.threadBuffer.get().getFloat());
+        return new FloatNode(this.buffer.getFloat());
     }
 
     private static BooleanNode decodeBoolean(int size)
@@ -317,7 +319,7 @@ private int[] sizeFromCtrlByte(int ctrlByte, int offset) {
     }
 
     private byte[] getByteArray(int length) {
-        return Decoder.getByteArray(this.threadBuffer.get(), length);
+        return Decoder.getByteArray(this.buffer, length);
     }
 
     private static byte[] getByteArray(ByteBuffer buffer, int length) {

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

@@ -19,10 +19,9 @@ public final class Reader implements Closeable {
             (byte) 0xCD, (byte) 0xEF, 'M', 'a', 'x', 'M', 'i', 'n', 'd', '.',
             'c', 'o', 'm' };
 
-    private int ipV4Start;
-    private final Decoder decoder;
+    private final int ipV4Start;
     private final Metadata metadata;
-    private final ThreadBuffer threadBuffer;
+    private final BufferHolder bufferHolder;
 
     /**
      * The file mode to use when opening a MaxMind DB.
@@ -63,7 +62,7 @@ public Reader(File database) throws IOException {
      *             if there is an error reading from the Stream.
      */
     public Reader(InputStream source) throws IOException {
-        this(new ThreadBuffer(source), "<InputStream>");
+        this(new BufferHolder(source), "<InputStream>");
     }
 
     /**
@@ -78,17 +77,19 @@ public Reader(InputStream source) throws IOException {
      *             if there is an error opening or reading from the file.
      */
     public Reader(File database, FileMode fileMode) throws IOException {
-        this(new ThreadBuffer(database, fileMode), database.getName());
+        this(new BufferHolder(database, fileMode), database.getName());
     }
 
-    private Reader(ThreadBuffer buffer, String name) throws IOException {
-        this.threadBuffer = buffer;
-        int start = this.findMetadataStart(name);
+    private Reader(BufferHolder bufferHolder, String name) throws IOException {
+        this.bufferHolder = bufferHolder;
 
-        Decoder metadataDecoder = new Decoder(this.threadBuffer, start);
+        ByteBuffer buffer = this.bufferHolder.get();
+        int start = this.findMetadataStart(buffer, name);
+
+        Decoder metadataDecoder = new Decoder(buffer, start);
         this.metadata = new Metadata(metadataDecoder.decode(start).getNode());
-        this.decoder = new Decoder(this.threadBuffer,
-                this.metadata.searchTreeSize + DATA_SECTION_SEPARATOR_SIZE);
+
+        this.ipV4Start = this.findIpV4StartNode(buffer);
     }
 
     /**
@@ -101,27 +102,28 @@ private Reader(ThreadBuffer buffer, String name) throws IOException {
      *             if a file I/O error occurs.
      */
     public JsonNode get(InetAddress ipAddress) throws IOException {
-        int pointer = this.findAddressInTree(ipAddress);
+        ByteBuffer buffer = this.bufferHolder.get();
+        int pointer = this.findAddressInTree(buffer, ipAddress);
         if (pointer == 0) {
             return null;
         }
-        return this.resolveDataPointer(pointer);
+        return this.resolveDataPointer(buffer, pointer);
     }
 
-    private int findAddressInTree(InetAddress address)
+    private int findAddressInTree(ByteBuffer buffer, InetAddress address)
             throws InvalidDatabaseException {
         byte[] rawAddress = address.getAddress();
 
         int bitLength = rawAddress.length * 8;
-        int record = this.startNode(bitLength);
+        int record = this.startNode(buffer, bitLength);
 
         for (int i = 0; i < bitLength; i++) {
             if (record >= this.metadata.nodeCount) {
                 break;
             }
             int b = 0xFF & rawAddress[i / 8];
             int bit = 1 & (b >> 7 - (i % 8));
-            record = this.readNode(record, bit);
+            record = this.readNode(buffer, record, bit);
         }
         if (record == this.metadata.nodeCount) {
             // record is empty
@@ -133,38 +135,31 @@ record = this.readNode(record, bit);
         throw new InvalidDatabaseException("Something bad happened");
     }
 
-    private int startNode(int bitLength) throws InvalidDatabaseException {
+    private int startNode(ByteBuffer buffer, int bitLength) throws InvalidDatabaseException {
         // 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.ipVersion == 6 && bitLength == 32) {
-            return this.ipV4StartNode();
+            return this.ipV4Start;
         }
         // The first node of the tree is always node 0, at the beginning of the
         // value
         return 0;
     }
 
-    private int ipV4StartNode() throws InvalidDatabaseException {
-        // This is a defensive check. There is no reason to call this when you
-        // have an IPv4 tree.
+    private int findIpV4StartNode(ByteBuffer buffer) throws InvalidDatabaseException {
         if (this.metadata.ipVersion == 4) {
             return 0;
         }
 
-        if (this.ipV4Start != 0) {
-            return this.ipV4Start;
-        }
         int node = 0;
         for (int i = 0; i < 96 && node < this.metadata.nodeCount; i++) {
-            node = this.readNode(node, 0);
+            node = this.readNode(buffer, node, 0);
         }
-        this.ipV4Start = node;
         return node;
     }
 
-    private int readNode(int nodeNumber, int index)
+    private int readNode(ByteBuffer buffer, int nodeNumber, int index)
             throws InvalidDatabaseException {
-        ByteBuffer buffer = this.threadBuffer.get();
         int baseOffset = nodeNumber * this.metadata.nodeByteSize;
 
         switch (this.metadata.recordSize) {
@@ -190,19 +185,21 @@ private int readNode(int nodeNumber, int index)
         }
     }
 
-    private JsonNode resolveDataPointer(int pointer) throws IOException {
+    private JsonNode resolveDataPointer(ByteBuffer buffer, int pointer) throws IOException {
         int resolved = (pointer - this.metadata.nodeCount)
                 + this.metadata.searchTreeSize;
 
-        if (resolved >= this.threadBuffer.get().capacity()) {
+        if (resolved >= buffer.capacity()) {
             throw new InvalidDatabaseException(
                     "The MaxMind DB file's search tree is corrupt: "
                             + "contains pointer larger than the database.");
         }
 
         // We only want the data from the decoder, not the offset where it was
         // found.
-        return this.decoder.decode(resolved).getNode();
+        Decoder decoder = new Decoder(buffer,
+                this.metadata.searchTreeSize + DATA_SECTION_SEPARATOR_SIZE);
+        return decoder.decode(resolved).getNode();
     }
 
     /*
@@ -213,9 +210,8 @@ private JsonNode resolveDataPointer(int pointer) throws IOException {
      * are much faster algorithms (e.g., Boyer-Moore) for this if speed is ever
      * an issue, but I suspect it won't be.
      */
-    private int findMetadataStart(String databaseName)
+    private int findMetadataStart(ByteBuffer buffer, String databaseName)
             throws InvalidDatabaseException {
-        ByteBuffer buffer = this.threadBuffer.get();
         int fileSize = buffer.capacity();
 
         FILE: for (int i = 0; i < fileSize - METADATA_START_MARKER.length + 1; i++) {
@@ -245,6 +241,6 @@ Metadata getMetadata() {
      */
     @Override
     public void close() throws IOException {
-        this.threadBuffer.close();
+        this.bufferHolder.close();
     }
 }

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

@@ -22,9 +22,6 @@
 import com.fasterxml.jackson.databind.node.ArrayNode;
 import com.fasterxml.jackson.databind.node.FloatNode;
 import com.fasterxml.jackson.databind.node.ObjectNode;
-import com.maxmind.db.Decoder;
-import com.maxmind.db.InvalidDatabaseException;
-import com.maxmind.db.ThreadBuffer;
 
 @SuppressWarnings({ "boxing", "static-method" })
 public class DecoderTest {
@@ -426,7 +423,7 @@ static <T> void testTypeDecoding(Decoder.Type type, Map<T, byte[]> tests)
             MappedByteBuffer mmap = fc.map(MapMode.READ_ONLY, 0, fc.size());
             try {
 
-                Decoder decoder = new Decoder(new ThreadBuffer(mmap), 0);
+                Decoder decoder = new Decoder(mmap, 0);
                 decoder.POINTER_TEST_HACK = true;
 
                 // XXX - this could be streamlined

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

@@ -10,9 +10,6 @@
 
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.databind.node.ObjectNode;
-import com.maxmind.db.Decoder;
-import com.maxmind.db.InvalidDatabaseException;
-import com.maxmind.db.ThreadBuffer;
 import com.maxmind.db.Reader.FileMode;
 
 public class PointerTest {
@@ -22,9 +19,9 @@ public void testWithPointers() throws InvalidDatabaseException,
             IOException, URISyntaxException {
         File file = new File(PointerTest.class.getResource(
                 "/maxmind-db/test-data/maps-with-pointers.raw").toURI());
-        ThreadBuffer ptf = new ThreadBuffer(file, FileMode.MEMORY);
+        BufferHolder ptf = new BufferHolder(file, FileMode.MEMORY);
         try {
-            Decoder decoder = new Decoder(ptf, 0);
+            Decoder decoder = new Decoder(ptf.get(), 0);
 
             ObjectMapper om = new ObjectMapper();