cache disk blocks in an LRU cache

Improves read access by factor 4 for small trees.

block-storage/src/main/java/org/lucares/pdb/blockstorage/BSFileDiskBlock.java

@@ -1,6 +1,6 @@
 package org.lucares.pdb.blockstorage;
 
-import java.nio.MappedByteBuffer;
+import java.nio.ByteBuffer;
 
 import org.lucares.collections.LongList;
 import org.lucares.pdb.diskstorage.DiskBlock;
@@ -28,7 +28,7 @@ public class BSFileDiskBlock {
 	public byte[] getBuffer() {
 
 		if (buffer == null) {
-			final MappedByteBuffer byteBuffer = diskBlock.getByteBuffer();
+			final ByteBuffer byteBuffer = diskBlock.getByteBuffer();
 			this.buffer = new byte[byteBuffer.capacity() - INT_SEQUENCE_OFFSET];
 			byteBuffer.position(INT_SEQUENCE_OFFSET);
 			byteBuffer.get(buffer);
@@ -86,7 +86,7 @@ public class BSFileDiskBlock {
 	}
 
 	public void force() {
-		diskBlock.getByteBuffer().force();
+		diskBlock.force();
 	}
 
 	@Override

block-storage/src/main/java/org/lucares/pdb/diskstorage/DiskBlock.java

@@ -1,5 +1,6 @@
 package org.lucares.pdb.diskstorage;
 
+import java.nio.ByteBuffer;
 import java.nio.MappedByteBuffer;
 
 public class DiskBlock {
@@ -7,9 +8,9 @@ public class DiskBlock {
 	private byte[] buffer = null;
 	private final long blockOffset;
 
-	private final MappedByteBuffer byteBuffer;
+	private final ByteBuffer byteBuffer;
 
-	public DiskBlock(final long blockOffset, final MappedByteBuffer byteBuffer) {
+	public DiskBlock(final long blockOffset, final ByteBuffer byteBuffer) {
 		this.blockOffset = blockOffset;
 		this.byteBuffer = byteBuffer;
 	}
@@ -24,7 +25,7 @@ public class DiskBlock {
 		return buffer;
 	}
 
-	public MappedByteBuffer getByteBuffer() {
+	public ByteBuffer getByteBuffer() {
 		return byteBuffer;
 	}
 
@@ -42,7 +43,10 @@ public class DiskBlock {
 	}
 
 	public void force() {
-		byteBuffer.force();
+		// some tests use HeapByteBuffer and don't support force
+		if (byteBuffer instanceof MappedByteBuffer) {
+			((MappedByteBuffer) byteBuffer).force();
+		}
 	}
 
 	@Override

block-storage/src/main/java/org/lucares/pdb/map/PersistentMap.java

@@ -16,6 +16,7 @@ import org.lucares.pdb.diskstorage.DiskBlock;
 import org.lucares.pdb.diskstorage.DiskStorage;
 import org.lucares.utils.Preconditions;
 import org.lucares.utils.byteencoder.VariableByteEncoder;
+import org.lucares.utils.cache.LRUCache;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -101,6 +102,8 @@ public class PersistentMap<K, V> implements AutoCloseable {
 
 	private final EncoderDecoder<V> valueEncoder;
 
+	private final LRUCache<Long, PersistentMapDiskNode> nodeCache = new LRUCache<>(10_000);
+
 	public PersistentMap(final Path path, final EncoderDecoder<K> keyEncoder, final EncoderDecoder<V> valueEncoder)
 			throws IOException {
 		this.diskStore = new DiskStorage(path);
@@ -294,16 +297,27 @@ public class PersistentMap<K, V> implements AutoCloseable {
 	}
 
 	private PersistentMapDiskNode getNode(final long nodeOffset) throws IOException {
-		final DiskBlock diskBlock = diskStore.getDiskBlock(nodeOffset, BLOCK_SIZE);
+
-		final byte[] buffer = diskBlock.getBuffer();
+		PersistentMapDiskNode node = nodeCache.get(nodeOffset);
-		final PersistentMapDiskNode node = PersistentMapDiskNode.parse(nodeOffset, buffer);
+		if (node == null) {
+
+			final DiskBlock diskBlock = diskStore.getDiskBlock(nodeOffset, BLOCK_SIZE);
+
+			node = PersistentMapDiskNode.parse(nodeOffset, diskBlock);
+			nodeCache.put(nodeOffset, node);
+		}
+
 		return node;
 	}
 
 	private void writeNode(final PersistentMapDiskNode node) throws IOException {
 		LOGGER.trace("writing node {}", node);
 		final long nodeOffest = node.getNodeOffset();
-		final DiskBlock diskBlock = diskStore.getDiskBlock(nodeOffest, BLOCK_SIZE);
+		// final DiskBlock diskBlock = diskStore.getDiskBlock(nodeOffest, BLOCK_SIZE);
+		DiskBlock diskBlock = node.getDiskBlock();
+		if (diskBlock == null) {
+			diskBlock = diskStore.getDiskBlock(nodeOffest, BLOCK_SIZE);
+		}
 		final byte[] buffer = diskBlock.getBuffer();
 		final byte[] newBuffer = node.serialize();
 		System.arraycopy(newBuffer, 0, buffer, 0, buffer.length);

block-storage/src/main/java/org/lucares/pdb/map/PersistentMapDiskNode.java

@@ -7,6 +7,7 @@ import java.util.List;
 import java.util.stream.Collectors;
 
 import org.lucares.collections.LongList;
+import org.lucares.pdb.diskstorage.DiskBlock;
 import org.lucares.pdb.map.NodeEntry.ValueType;
 import org.lucares.utils.Preconditions;
 import org.lucares.utils.byteencoder.VariableByteEncoder;
@@ -32,19 +33,23 @@ import org.lucares.utils.byteencoder.VariableByteEncoder;
  */
 public class PersistentMapDiskNode {
 
+	// TODO use map instead of list
 	private final List<NodeEntry> entries;
 	private final long nodeOffset;
+	private final DiskBlock diskBlock;
 
-	public PersistentMapDiskNode(final long nodeOffset, final List<NodeEntry> entries) {
+	public PersistentMapDiskNode(final long nodeOffset, final List<NodeEntry> entries, final DiskBlock diskBlock) {
 		this.nodeOffset = nodeOffset;
+		this.diskBlock = diskBlock;
 		this.entries = new ArrayList<>(entries);
 	}
 
 	public static PersistentMapDiskNode emptyRootNode(final long nodeOffset) {
-		return new PersistentMapDiskNode(nodeOffset, Collections.emptyList());
+		return new PersistentMapDiskNode(nodeOffset, Collections.emptyList(), null);
 	}
 
-	public static PersistentMapDiskNode parse(final long nodeOffset, final byte[] data) {
+	public static PersistentMapDiskNode parse(final long nodeOffset, final DiskBlock diskBlock) {
+		final byte[] data = diskBlock.getBuffer();
 		if (data.length != PersistentMap.BLOCK_SIZE) {
 			throw new IllegalStateException(
 					"block size must be " + PersistentMap.BLOCK_SIZE + " but was " + data.length);
@@ -52,7 +57,7 @@ public class PersistentMapDiskNode {
 		final LongList longs = VariableByteEncoder.decode(data);
 
 		final List<NodeEntry> entries = deserialize(longs, data);
-		return new PersistentMapDiskNode(nodeOffset, entries);
+		return new PersistentMapDiskNode(nodeOffset, entries, diskBlock);
 	}
 
 	public static List<NodeEntry> deserialize(final LongList keyLengths, final byte[] buffer) {
@@ -91,6 +96,10 @@ public class PersistentMapDiskNode {
 		return serialize(entries);
 	}
 
+	public DiskBlock getDiskBlock() {
+		return diskBlock;
+	}
+
 	public long getNodeOffset() {
 		return nodeOffset;
 	}
@@ -195,7 +204,7 @@ public class PersistentMapDiskNode {
 		entries.clear();
 		entries.addAll(rightEntries);
 
-		return new PersistentMapDiskNode(newBlockOffset, leftEntries);
+		return new PersistentMapDiskNode(newBlockOffset, leftEntries, null);
 	}
 
 	public static int neededBytesTotal(final List<NodeEntry> entries) {

block-storage/src/test/java/org/lucares/pdb/map/PersistentMapDiskNodeTest.java

@@ -1,10 +1,12 @@
 package org.lucares.pdb.map;
 
+import java.nio.ByteBuffer;
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.ThreadLocalRandom;
 
+import org.lucares.pdb.diskstorage.DiskBlock;
 import org.lucares.pdb.map.NodeEntry.ValueType;
 import org.testng.Assert;
 import org.testng.annotations.Test;
@@ -21,11 +23,13 @@ public class PersistentMapDiskNodeTest {
 		entries.add(newNode(ValueType.VALUE_INLINE, "key4___", "value4----"));
 
 		final long nodeOffset = ThreadLocalRandom.current().nextInt();
-		final PersistentMapDiskNode node = new PersistentMapDiskNode(nodeOffset, entries);
+		final PersistentMapDiskNode node = new PersistentMapDiskNode(nodeOffset, entries, null);
 
 		final byte[] buffer = node.serialize();
 
-		final PersistentMapDiskNode actualNode = PersistentMapDiskNode.parse(nodeOffset, buffer);
+		final ByteBuffer byteBuffer = ByteBuffer.wrap(buffer);
+		final PersistentMapDiskNode actualNode = PersistentMapDiskNode.parse(nodeOffset,
+				new DiskBlock(nodeOffset, byteBuffer));
 
 		Assert.assertEquals(actualNode.getEntries(), entries);
 	}

block-storage/src/test/java/org/lucares/pdb/map/PersistentMapTest.java

@@ -57,6 +57,7 @@ public class PersistentMapTest {
 		}
 	}
 
+	@Test(invocationCount = 1)
 	public void testManyValues() throws Exception {
 		final Path file = dataDirectory.resolve("map.db");
 		final var insertedValues = new HashMap<String, String>();
@@ -118,7 +119,7 @@ public class PersistentMapTest {
 		}
 	}
 
-	@Test
+	@Test(invocationCount = 1)
 	public void testManySmallValues() throws Exception {
 		final Path file = dataDirectory.resolve("map.db");
 		final var insertedValues = new HashMap<Long, Long>();
@@ -267,4 +268,64 @@ public class PersistentMapTest {
 		}
 	}
 
+	@Test(invocationCount = 1)
+	public void testLotsOfValues() throws Exception {
+		final Path file = dataDirectory.resolve("map.db");
+		final var insertedValues = new HashMap<Long, Long>();
+
+		final SecureRandom rnd = new SecureRandom();
+		rnd.setSeed(1);
+
+		try (final PersistentMap<Long, Long> map = new PersistentMap<>(file, PersistentMap.LONG_CODER,
+				PersistentMap.LONG_CODER)) {
+
+			for (int i = 0; i < 1_000; i++) {
+
+				final Long key = (long) (rnd.nextGaussian() * Integer.MAX_VALUE);
+				final Long value = (long) (rnd.nextGaussian() * Integer.MAX_VALUE);
+
+				if (insertedValues.containsKey(key)) {
+					continue;
+				}
+
+				Assert.assertNull(map.putValue(key, value));
+
+				insertedValues.put(key, value);
+
+				final boolean failEarly = false;
+				if (failEarly) {
+					for (final var entry : insertedValues.entrySet()) {
+						final Long actualValue = map.getValue(entry.getKey());
+
+						if (!Objects.equals(actualValue, entry.getValue())) {
+							map.print();
+						}
+
+						Assert.assertEquals(actualValue, entry.getValue(),
+								"value for key " + entry.getKey() + " in the " + i + "th iteration");
+					}
+				}
+			}
+		}
+
+		try (final PersistentMap<Long, Long> map = new PersistentMap<>(file, PersistentMap.LONG_CODER,
+				PersistentMap.LONG_CODER)) {
+			final AtomicInteger counter = new AtomicInteger();
+			final AtomicInteger maxDepth = new AtomicInteger();
+			map.visitNodeEntriesPreOrder((node, parentNode, nodeEntry, depth) -> {
+				counter.addAndGet(nodeEntry.isInnerNode() ? 1 : 0);
+				maxDepth.set(Math.max(maxDepth.get(), depth));
+			});
+
+			final long start = System.nanoTime();
+			for (final var entry : insertedValues.entrySet()) {
+				final Long actualValue = map.getValue(entry.getKey());
+				Assert.assertEquals(actualValue, entry.getValue(),
+						"value for key " + entry.getKey() + " after all iterations");
+			}
+			System.out.println("nodes=" + counter.get() + ", depth=" + maxDepth.get() + ": "
+					+ (System.nanoTime() - start) / 1_000_000.0 + "ms");
+		}
+	}
+
 }

pdb-utils/src/main/java/org/lucares/utils/cache/LRUCache.java

@@ -0,0 +1,35 @@
+package org.lucares.utils.cache;
+
+import java.util.LinkedHashMap;
+import java.util.Map;
+
+public class LRUCache<K, V> {
+	private final LinkedHashMap<K, V> cache;
+
+	public LRUCache(final int maxEntries) {
+		this.cache = new LinkedHashMap<>(16, 0.75f, true) {
+			private static final long serialVersionUID = 1L;
+
+			protected boolean removeEldestEntry(final Map.Entry<K, V> eldest) {
+				return size() > maxEntries;
+			}
+		};
+	}
+
+	public V put(final K key, final V value) {
+		return cache.put(key, value);
+	}
+
+	public V get(final K key) {
+		return cache.get(key);
+	}
+
+	public V remove(final K key) {
+		return cache.remove(key);
+	}
+
+	public int size() {
+		return cache.size();
+	}
+
+}