From 10d4ac600dce82c836f6a0db6073f57e74620fca Mon Sep 17 00:00:00 2001
From: Andrea Cavalli <andrea@cavallium.it>
Date: Thu, 5 Jan 2023 02:21:09 +0100
Subject: [PATCH] More precise statistics

---
 .../database/disk/AbstractRocksDBColumn.java     | 13 +++++++++++++
 .../database/disk/LLLocalDictionary.java         | 16 ++++++++++++----
 .../dbengine/database/disk/RocksDBColumn.java    |  2 ++
 3 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/src/main/java/it/cavallium/dbengine/database/disk/AbstractRocksDBColumn.java b/src/main/java/it/cavallium/dbengine/database/disk/AbstractRocksDBColumn.java
index 11c334c..2603342 100644
--- a/src/main/java/it/cavallium/dbengine/database/disk/AbstractRocksDBColumn.java
+++ b/src/main/java/it/cavallium/dbengine/database/disk/AbstractRocksDBColumn.java
@@ -26,6 +26,7 @@ import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.nio.charset.StandardCharsets;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.locks.StampedLock;
 import org.apache.logging.log4j.LogManager;
@@ -41,11 +42,13 @@ import org.rocksdb.FlushOptions;
 import org.rocksdb.Holder;
 import org.rocksdb.KeyMayExist;
 import org.rocksdb.KeyMayExist.KeyMayExistEnum;
+import org.rocksdb.Range;
 import org.rocksdb.ReadOptions;
 import org.rocksdb.RocksDB;
 import org.rocksdb.RocksDBException;
 import org.rocksdb.RocksObject;
 import org.rocksdb.Slice;
+import org.rocksdb.TableProperties;
 import org.rocksdb.Transaction;
 import org.rocksdb.TransactionOptions;
 import org.rocksdb.WriteBatch;
@@ -653,6 +656,16 @@ public sealed abstract class AbstractRocksDBColumn<T extends RocksDB> implements
 		}
 	}
 
+	@Override
+	public long getNumEntries() throws RocksDBException {
+		Map<String, TableProperties> props = db.getPropertiesOfAllTables(cfh);
+		long entries = 0;
+		for (TableProperties tableProperties : props.values()) {
+			entries += tableProperties.getNumEntries();
+		}
+		return entries;
+	}
+
 	@Override
 	public void write(WriteOptions writeOptions, WriteBatch writeBatch) throws RocksDBException {
 		var closeReadLock = closeLock.readLock();
diff --git a/src/main/java/it/cavallium/dbengine/database/disk/LLLocalDictionary.java b/src/main/java/it/cavallium/dbengine/database/disk/LLLocalDictionary.java
index 52fa7df..7a8afa8 100644
--- a/src/main/java/it/cavallium/dbengine/database/disk/LLLocalDictionary.java
+++ b/src/main/java/it/cavallium/dbengine/database/disk/LLLocalDictionary.java
@@ -27,7 +27,6 @@ import it.cavallium.dbengine.database.LLRange;
 import it.cavallium.dbengine.database.LLSnapshot;
 import it.cavallium.dbengine.database.LLUtils;
 import it.cavallium.dbengine.database.OptionalBuf;
-import it.cavallium.dbengine.database.SafeCloseable;
 import it.cavallium.dbengine.database.UpdateMode;
 import it.cavallium.dbengine.database.UpdateReturnMode;
 import it.cavallium.dbengine.database.serialization.KVSerializationFunction;
@@ -35,10 +34,8 @@ import it.cavallium.dbengine.rpc.current.data.DatabaseOptions;
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.time.Duration;
-import java.time.Instant;
 import java.util.ArrayList;
 import java.util.List;
-import java.util.Optional;
 import java.util.concurrent.Callable;
 import java.util.concurrent.ForkJoinPool;
 import java.util.concurrent.ForkJoinTask;
@@ -104,6 +101,7 @@ public class LLLocalDictionary implements LLDictionary {
 	 */
 	static final boolean USE_WRITE_BATCH_IN_SET_RANGE_DELETE = false;
 	static final boolean PARALLEL_EXACT_SIZE = true;
+	private static final boolean USE_NUM_ENTRIES_PRECISE_COUNTER = true;
 
 	private static final byte[] FIRST_KEY = new byte[]{};
 	/**
@@ -1253,12 +1251,15 @@ public class LLLocalDictionary implements LLDictionary {
 		try (var rocksdbSnapshot = generateReadOptionsOrNew(snapshot)) {
 			if (USE_CURRENT_FASTSIZE_FOR_OLD_SNAPSHOTS || rocksdbSnapshot.snapshot() == null) {
 				try {
+					if (USE_NUM_ENTRIES_PRECISE_COUNTER) {
+						return exactSizeAll(null);
+					}
 					return db.getLongProperty("rocksdb.estimate-num-keys");
 				} catch (RocksDBException e) {
 					logger.error(MARKER_ROCKSDB, "Failed to get RocksDB estimated keys count property", e);
 					return 0;
 				}
-			} else if (PARALLEL_EXACT_SIZE) {
+			} else if (USE_NUM_ENTRIES_PRECISE_COUNTER || PARALLEL_EXACT_SIZE) {
 				return exactSizeAll(snapshot);
 			} else {
 				rocksdbSnapshot.setFillCache(false);
@@ -1282,6 +1283,13 @@ public class LLLocalDictionary implements LLDictionary {
 		if (Schedulers.isInNonBlockingThread()) {
 			throw new UnsupportedOperationException("Called exactSizeAll in a nonblocking thread");
 		}
+		if (snapshot == null && USE_NUM_ENTRIES_PRECISE_COUNTER) {
+			try {
+				return db.getNumEntries();
+			} catch (RocksDBException ex) {
+				throw new IllegalStateException("Failed to read exact size", ex);
+			}
+		}
 		try (var readOpts = LLUtils.generateCustomReadOptions(generateReadOptionsOrNull(snapshot), false, false, false)) {
 			if (LLUtils.MANUAL_READAHEAD) {
 				readOpts.setReadaheadSize(128 * 1024); // 128KiB
diff --git a/src/main/java/it/cavallium/dbengine/database/disk/RocksDBColumn.java b/src/main/java/it/cavallium/dbengine/database/disk/RocksDBColumn.java
index 4cab0d1..2521173 100644
--- a/src/main/java/it/cavallium/dbengine/database/disk/RocksDBColumn.java
+++ b/src/main/java/it/cavallium/dbengine/database/disk/RocksDBColumn.java
@@ -91,6 +91,8 @@ public sealed interface RocksDBColumn permits AbstractRocksDBColumn {
 
 	long getLongProperty(String property) throws RocksDBException;
 
+	long getNumEntries() throws RocksDBException;
+
 	ColumnFamilyHandle getColumnFamilyHandle();
 
 	BufferAllocator getAllocator();