Made the buffer thread local, which should solve most of our thread safety issues.

Author: oschwald

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

@@ -17,7 +17,7 @@
 import com.fasterxml.jackson.databind.node.ObjectNode;
 import com.fasterxml.jackson.databind.node.TextNode;
 
-public class Decoder {
+public final class Decoder {
 
     private final boolean DEBUG;
     // XXX - This is only for unit testings. We should possibly make a
@@ -27,7 +27,27 @@ public class Decoder {
 
     private final ObjectMapper objectMapper;
 
-    private final ByteBuffer mmap;
+    private final ThreadBuffer threadBuffer;
+
+    static enum Type {
+        EXTENDED, POINTER, UTF8_STRING, DOUBLE, BYTES, UINT16, UINT32, MAP, INT32, UINT64, UINT128, ARRAY, CONTAINER, END_MARKER, BOOLEAN;
+
+        public static Type get(int i) {
+            // XXX - Type.values() might be expensive. Consider caching it.
+            return Type.values()[i];
+        }
+
+        private static Type get(byte b) {
+            // bytes are signed, but we want to treat them as unsigned here
+            // XXX - Type.values() might be expensive. Consider caching it.
+            return Type.get(b & 0xFF);
+        }
+
+        public static Type fromControlByte(int b) {
+            // The type is encoded in the first 3 bits of the byte.
+            return Type.get((byte) ((0xFF & b) >>> 5));
+        }
+    }
 
     class Result {
         private final JsonNode node;
@@ -52,21 +72,23 @@ void setOffset(long offset) {
 
     }
 
-    public Decoder(ByteBuffer mmap, long pointerBase) {
+    public Decoder(ThreadBuffer threadBuffer, long pointerBase) {
         this.pointerBase = pointerBase;
-        this.mmap = mmap;
+        this.threadBuffer = threadBuffer;
         this.objectMapper = new ObjectMapper();
         this.DEBUG = System.getenv().get("MAXMIND_DB_DECODER_DEBUG") != null;
     }
 
     public Result decode(long offset) throws MaxMindDbException, IOException {
-        this.mmap.position((int) offset);
+        ByteBuffer buffer = this.threadBuffer.get();
+
+        buffer.position((int) offset);
 
         if (this.DEBUG) {
             Log.debug("Offset", String.valueOf(offset));
         }
 
-        int ctrlByte = 0xFF & this.mmap.get();
+        int ctrlByte = 0xFF & buffer.get();
 
         offset++;
 
@@ -95,7 +117,7 @@ public Result decode(long offset) throws MaxMindDbException, IOException {
         }
 
         if (type.equals(Type.EXTENDED)) {
-            int nextByte = this.mmap.get();
+            int nextByte = buffer.get();
 
             if (this.DEBUG) {
                 Log.debug("Next byte", nextByte);
@@ -210,7 +232,7 @@ private Result decodePointer(int ctrlByte, long offset) {
     }
 
     private String decodeString(int size) {
-        ByteBuffer buffer = this.mmap.slice();
+        ByteBuffer buffer = this.threadBuffer.get().slice();
         buffer.limit(size);
         return Charset.forName("UTF-8").decode(buffer).toString();
     }
@@ -228,9 +250,10 @@ private int decodeInteger(int size) {
     }
 
     private int decodeInteger(int base, int size) {
+        ByteBuffer buffer = this.threadBuffer.get();
         int integer = base;
         for (int i = 0; i < size; i++) {
-            integer = (integer << 8) | (this.mmap.get() & 0xFF);
+            integer = (integer << 8) | (buffer.get() & 0xFF);
         }
         return integer;
     }
@@ -244,9 +267,10 @@ long decodeLong(int size) {
     }
 
     long decodeLong(long base, int size) {
+        ByteBuffer buffer = this.threadBuffer.get();
         long longInt = base;
         for (int i = 0; i < size; i++) {
-            longInt = (longInt << 8) | (this.mmap.get() & 0xFF);
+            longInt = (longInt << 8) | (buffer.get() & 0xFF);
         }
         return longInt;
     }
@@ -325,8 +349,7 @@ private Result decodeMap(long size, long offset) throws MaxMindDbException,
 
     }
 
-    private long[] sizeFromCtrlByte(int ctrlByte, long offset)
-            throws IOException {
+    private long[] sizeFromCtrlByte(int ctrlByte, long offset) {
         int size = ctrlByte & 0x1f;
 
         if (size < 29) {
@@ -351,6 +374,28 @@ private long[] sizeFromCtrlByte(int ctrlByte, long offset)
     }
 
     private byte[] getByteArray(int length) {
-        return Util.getByteArray(this.mmap, length);
+        return Decoder.getByteArray(this.threadBuffer.get(), length);
+    }
+
+    static int decodeInteger(byte[] bytes) {
+        int i = 0;
+        for (byte b : bytes) {
+            i = (i << 8) | (b & 0xFF);
+        }
+        return i;
+    }
+
+    static long decodeLong(byte[] bytes) {
+        long i = 0;
+        for (byte b : bytes) {
+            i = (i << 8) | (b & 0xFF);
+        }
+        return i;
+    }
+
+    static byte[] getByteArray(ByteBuffer buffer, int length) {
+        byte[] bytes = new byte[length];
+        buffer.get(bytes);
+        return bytes;
     }
 }

src/main/java/com/maxmind/maxminddb/Log.java

@@ -3,7 +3,7 @@
 import java.net.InetAddress;
 import java.nio.ByteBuffer;
 
-class Log {
+final class Log {
 
     static void debug(String string, ByteBuffer buffer) {
         debug(string, buffer.array());

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

@@ -5,7 +5,7 @@
 import com.fasterxml.jackson.databind.JsonNode;
 
 // XXX - if we make this public, add getters
-class Metadata {
+final class Metadata {
     final int binaryFormatMajorVersion;
     final int binaryFormatMinorVersion;
     private final BigInteger buildEpoch;

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

@@ -2,17 +2,13 @@
 
 import java.io.File;
 import java.io.IOException;
-import java.io.RandomAccessFile;
 import java.net.InetAddress;
 import java.nio.ByteBuffer;
-import java.nio.MappedByteBuffer;
-import java.nio.channels.FileChannel;
-import java.nio.channels.FileChannel.MapMode;
 import java.util.Arrays;
 
 import com.fasterxml.jackson.databind.JsonNode;
 
-public class Reader {
+public final class Reader {
     private static final int DATA_SECTION_SEPARATOR_SIZE = 16;
     private static final byte[] METADATA_START_MARKER = { (byte) 0xAB,
             (byte) 0xCD, (byte) 0xEF, 'M', 'a', 'x', 'M', 'i', 'n', 'd', '.',
@@ -21,15 +17,23 @@ public class Reader {
     private final boolean DEBUG;
     private final Decoder decoder;
     private final Metadata metadata;
-    private final FileChannel fc;
-    private final RandomAccessFile raf;
-    private final MappedByteBuffer mmap;
+    private final ThreadBuffer threadBuffer;
+
+    public enum FileMode {
+        MEMORY_MAPPED, IN_MEMORY
+    }
 
     public Reader(File database) throws MaxMindDbException, IOException {
+        this(database, FileMode.MEMORY_MAPPED);
+    }
+
+    // XXX - loading the file into memory doesn't really provide any performance
+    // gains on my machine. Consider whether it is even worth providing the
+    // option.
+    public Reader(File database, FileMode mode) throws MaxMindDbException,
+            IOException {
         this.DEBUG = System.getenv().get("MAXMIND_DB_READER_DEBUG") != null;
-        this.raf = new RandomAccessFile(database, "r");
-        this.fc = this.raf.getChannel();
-        this.mmap = this.fc.map(MapMode.READ_ONLY, 0, this.fc.size());
+        this.threadBuffer = new ThreadBuffer(database, mode);
 
         /*
          * We need to make sure that whatever chunk we read will have the
@@ -54,12 +58,12 @@ public Reader(File database) throws MaxMindDbException, IOException {
                             + "). Is this a valid MaxMind DB file?");
         }
 
-        Decoder metadataDecoder = new Decoder(this.mmap, 0);
+        Decoder metadataDecoder = new Decoder(this.threadBuffer, 0);
 
         this.metadata = new Metadata(metadataDecoder.decode(start).getNode());
 
-        this.decoder = new Decoder(this.mmap, this.metadata.searchTreeSize
-                + DATA_SECTION_SEPARATOR_SIZE);
+        this.decoder = new Decoder(this.threadBuffer,
+                this.metadata.searchTreeSize + DATA_SECTION_SEPARATOR_SIZE);
 
         if (this.DEBUG) {
             Log.debug(this.metadata.toString());
@@ -75,12 +79,11 @@ public JsonNode get(InetAddress address) throws MaxMindDbException,
             return null;
         }
 
-        this.mmap.position((int) pointer);
+        this.threadBuffer.get().position((int) pointer);
         return this.resolveDataPointer(pointer);
     }
 
-    long findAddressInTree(InetAddress address) throws MaxMindDbException,
-            IOException {
+    long findAddressInTree(InetAddress address) throws MaxMindDbException {
         byte[] rawAddress = address.getAddress();
 
         // XXX sort of wasteful
@@ -141,35 +144,34 @@ long record = nodes[bit];
         throw new MaxMindDbException("Something bad happened");
     }
 
-    private long[] readNode(long nodeNumber) throws IOException,
-            MaxMindDbException {
-        this.mmap.position((int) nodeNumber * this.metadata.nodeByteSize);
+    private long[] readNode(long nodeNumber) throws MaxMindDbException {
+        ByteBuffer buffer = this.threadBuffer.get();
+        buffer.position((int) nodeNumber * this.metadata.nodeByteSize);
 
-        byte[] buffer = Util
-                .getByteArray(this.mmap, this.metadata.nodeByteSize);
+        byte[] bytes = Decoder.getByteArray(buffer, this.metadata.nodeByteSize);
 
         if (this.DEBUG) {
-            Log.debug("Node bytes", buffer);
+            Log.debug("Node bytes", bytes);
         }
-        return this.splitNodeIntoRecords(buffer);
+        return this.splitNodeIntoRecords(bytes);
     }
 
     private long[] splitNodeIntoRecords(byte[] bytes) throws MaxMindDbException {
         long[] nodes = new long[2];
         switch (this.metadata.recordSize) {
             case 24:
-                nodes[0] = Util.decodeLong(Arrays.copyOfRange(bytes, 0, 3));
-                nodes[1] = Util.decodeLong(Arrays.copyOfRange(bytes, 3, 6));
+                nodes[0] = Decoder.decodeLong(Arrays.copyOfRange(bytes, 0, 3));
+                nodes[1] = Decoder.decodeLong(Arrays.copyOfRange(bytes, 3, 6));
                 return nodes;
             case 28:
-                nodes[0] = Util.decodeLong(Arrays.copyOfRange(bytes, 0, 3));
-                nodes[1] = Util.decodeLong(Arrays.copyOfRange(bytes, 4, 7));
+                nodes[0] = Decoder.decodeLong(Arrays.copyOfRange(bytes, 0, 3));
+                nodes[1] = Decoder.decodeLong(Arrays.copyOfRange(bytes, 4, 7));
                 nodes[0] = ((0xF0 & bytes[3]) << 20) | nodes[0];
                 nodes[1] = ((0x0F & bytes[3]) << 24) | nodes[1];
                 return nodes;
             case 32:
-                nodes[0] = Util.decodeLong(Arrays.copyOfRange(bytes, 0, 4));
-                nodes[1] = Util.decodeLong(Arrays.copyOfRange(bytes, 4, 8));
+                nodes[0] = Decoder.decodeLong(Arrays.copyOfRange(bytes, 0, 4));
+                nodes[1] = Decoder.decodeLong(Arrays.copyOfRange(bytes, 4, 8));
                 return nodes;
             default:
                 throw new MaxMindDbException("Unknown record size: "
@@ -203,15 +205,15 @@ private JsonNode resolveDataPointer(long pointer)
      * are much faster algorithms (e.g., Boyer-Moore) for this if speed is ever
      * an issue, but I suspect it won't be.
      */
-    private long findMetadataStart() throws IOException {
-        long fileSize = this.fc.size();
+    private long findMetadataStart() {
+        ByteBuffer buffer = this.threadBuffer.get();
+        int fileSize = buffer.capacity();
 
-        FILE: for (long i = 0; i < fileSize - METADATA_START_MARKER.length + 1; i++) {
+        FILE: for (int i = 0; i < fileSize - METADATA_START_MARKER.length + 1; i++) {
             for (int j = 0; j < METADATA_START_MARKER.length; j++) {
-                ByteBuffer b = ByteBuffer.wrap(new byte[1]);
-                this.fc.read(b, fileSize - i - j - 1);
-                if (b.get(0) != METADATA_START_MARKER[METADATA_START_MARKER.length
-                        - j - 1]) {
+                byte b = buffer.get(fileSize - i - j - 1);
+                if (b != METADATA_START_MARKER[METADATA_START_MARKER.length - j
+                        - 1]) {
                     continue FILE;
                 }
             }
@@ -220,16 +222,11 @@ private long findMetadataStart() throws IOException {
         return -1;
     }
 
-    public Metadata getMetadata() {
+    Metadata getMetadata() {
         return this.metadata;
     }
 
     public void close() throws IOException {
-        if (this.fc != null) {
-            this.fc.close();
-        }
-        if (this.raf != null) {
-            this.raf.close();
-        }
+        this.threadBuffer.close();
     }
 }

src/main/java/com/maxmind/maxminddb/ThreadBuffer.java

@@ -0,0 +1,49 @@
+package com.maxmind.maxminddb;
+
+import java.io.File;
+import java.io.IOException;
+import java.io.RandomAccessFile;
+import java.nio.ByteBuffer;
+import java.nio.channels.FileChannel;
+import java.nio.channels.FileChannel.MapMode;
+
+import com.maxmind.maxminddb.Reader.FileMode;
+
+final class ThreadBuffer extends ThreadLocal<ByteBuffer> {
+    // XXX - DO NOT PASS THIS OUTSIDE THIS CLASS.
+    private final ByteBuffer buffer;
+    private final RandomAccessFile raf;
+    private final FileChannel fc;
+
+    ThreadBuffer(File database, FileMode mode) throws IOException {
+        this.raf = new RandomAccessFile(database, "r");
+        this.fc = this.raf.getChannel();
+        if (mode == FileMode.IN_MEMORY) {
+            this.buffer = ByteBuffer.wrap(new byte[(int) this.fc.size()]);
+            this.fc.read(this.buffer);
+        } else {
+            this.buffer = this.fc.map(MapMode.READ_ONLY, 0, this.fc.size());
+        }
+    }
+
+    // This is just to ease unit testing
+    ThreadBuffer(ByteBuffer buffer) {
+        this.buffer = buffer;
+        this.raf = null;
+        this.fc = null;
+    }
+
+    @Override
+    protected synchronized ByteBuffer initialValue() {
+        return this.buffer.duplicate();
+    }
+
+    public void close() throws IOException {
+        if (this.fc != null) {
+            this.fc.close();
+        }
+        if (this.raf != null) {
+            this.raf.close();
+        }
+    }
+}