Safer access to database elements

src/main/java/it/cavallium/dbengine/database/disk/AbstractRocksDBColumn.java

@@ -332,17 +332,11 @@ public sealed abstract class AbstractRocksDBColumn<T extends RocksDB> implements
 	}
 
 	protected void ensureOpen() {
-		if (Schedulers.isInNonBlockingThread()) {
+		RocksDBUtils.ensureOpen(db, cfh);
-			throw new UnsupportedOperationException("Called in a nonblocking thread");
-		}
-		ensureOwned(db);
-		ensureOwned(cfh);
 	}
 
 	protected void ensureOwned(org.rocksdb.RocksObject rocksObject) {
-		if (!rocksObject.isOwningHandle()) {
+		RocksDBUtils.ensureOwned(rocksObject);
-			throw new IllegalStateException("Not owning handle");
-		}
 	}
 
 	@Override
@@ -939,7 +933,13 @@ public sealed abstract class AbstractRocksDBColumn<T extends RocksDB> implements
 	}
 
 	protected int getLevels() {
-		return RocksDBUtils.getLevels(db, cfh);
+		var closeReadLock = closeLock.readLock();
+		try {
+			ensureOpen();
+			return RocksDBUtils.getLevels(db, cfh);
+		} finally {
+			closeLock.unlockRead(closeReadLock);
+		}
 	}
 
 	@Override

src/main/java/it/cavallium/dbengine/database/disk/LLLocalKeyValueDatabase.java

@@ -516,6 +516,17 @@ public class LLLocalKeyValueDatabase implements LLKeyValueDatabase {
 				|| Boolean.parseBoolean(System.getProperty("it.cavallium.dbengine.disableslowdown", "false"));
 	}
 
+	protected void ensureOpen() {
+		if (closed) {
+			throw new IllegalStateException("Database closed");
+		}
+		RocksDBUtils.ensureOpen(db, null);
+	}
+
+	protected void ensureOwned(org.rocksdb.RocksObject rocksObject) {
+		RocksDBUtils.ensureOwned(rocksObject);
+	}
+
 	private synchronized PersistentCache resolvePersistentCache(HashMap<String, PersistentCache> caches,
 			DBOptions rocksdbOptions,
 			List<it.cavallium.dbengine.rpc.current.data.PersistentCache> persistentCaches,
@@ -562,22 +573,51 @@ public class LLLocalKeyValueDatabase implements LLKeyValueDatabase {
 	}
 
 	public int getLevels(Column column) {
-		return RocksDBUtils.getLevels(db, handles.get(column));
+		var closeReadLock = closeLock.readLock();
+		try {
+			ensureOpen();
+			var cfh = handles.get(column);
+			ensureOwned(cfh);
+			return RocksDBUtils.getLevels(db, cfh);
+		} finally {
+			closeLock.unlockRead(closeReadLock);
+		}
 	}
 
 	public List<String> getColumnFiles(Column column, boolean excludeLastLevel) {
-		var cfh = handles.get(column);
+		var closeReadLock = closeLock.readLock();
-		return RocksDBUtils.getColumnFiles(db, cfh, excludeLastLevel);
+		try {
+			ensureOpen();
+			var cfh = handles.get(column);
+			ensureOwned(cfh);
+			return RocksDBUtils.getColumnFiles(db, cfh, excludeLastLevel);
+		} finally {
+			closeLock.unlockRead(closeReadLock);
+		}
 	}
 
 	public void forceCompaction(int volumeId) throws RocksDBException {
-		for (var cfh : this.handles.values()) {
+		var closeReadLock = closeLock.readLock();
-			RocksDBUtils.forceCompaction(db, name, cfh, volumeId, logger);
+		try {
+			ensureOpen();
+			for (var cfh : this.handles.values()) {
+				ensureOwned(cfh);
+				RocksDBUtils.forceCompaction(db, name, cfh, volumeId, logger);
+			}
+		} finally {
+			closeLock.unlockRead(closeReadLock);
 		}
 	}
 
 	public void flush(FlushOptions flushOptions) throws RocksDBException {
-		db.flush(flushOptions);
+		var closeReadLock = closeLock.readLock();
+		try {
+			ensureOpen();
+			ensureOwned(flushOptions);
+			db.flush(flushOptions);
+		} finally {
+			closeLock.unlockRead(closeReadLock);
+		}
 	}
 
 	private record RocksLevelOptions(CompressionType compressionType, CompressionOptions compressionOptions) {}
@@ -601,7 +641,11 @@ public class LLLocalKeyValueDatabase implements LLKeyValueDatabase {
 					if (closed) {
 						return 0d;
 					}
+					var closeReadLock = closeLock.readLock();
 					try {
+						if (closed) {
+							return 0d;
+						}
 						return database.getAggregatedLongProperty(propertyName)
 								/ (divideByAllColumns ? getAllColumnFamilyHandles().size() : 1d);
 					} catch (RocksDBException e) {
@@ -609,6 +653,8 @@ public class LLLocalKeyValueDatabase implements LLKeyValueDatabase {
 							return 0d;
 						}
 						throw new RuntimeException(e);
+					} finally {
+						closeLock.unlockRead(closeReadLock);
 					}
 				}
 		);
@@ -949,19 +995,41 @@ public class LLLocalKeyValueDatabase implements LLKeyValueDatabase {
 		}
 	}
 
+	private Snapshot getSnapshotLambda(LLSnapshot snapshot) {
+		var closeReadSnapLock = closeLock.readLock();
+		try {
+			ensureOpen();
+			var snapshotHandle = snapshotsHandles.get(snapshot.getSequenceNumber());
+			ensureOwned(snapshotHandle);
+			return snapshotHandle;
+		} finally {
+			closeLock.unlockRead(closeReadSnapLock);
+		}
+	}
+
 	@Override
 	public Mono<LLLocalSingleton> getSingleton(byte[] singletonListColumnName,
 			byte[] name,
 			byte @Nullable[] defaultValue) {
 		return Mono
-				.fromCallable(() -> new LLLocalSingleton(
+				.fromCallable(() -> {
-						getRocksDBColumn(db, getCfh(singletonListColumnName)),
+					var closeReadLock = closeLock.readLock();
-						(snapshot) -> snapshotsHandles.get(snapshot.getSequenceNumber()),
+					try {
-						LLLocalKeyValueDatabase.this.name,
+						ensureOpen();
-						name,
+						var cfh = getCfh(singletonListColumnName);
-						ColumnUtils.toString(singletonListColumnName),
+						ensureOwned(cfh);
-						dbWScheduler, dbRScheduler, defaultValue
+						return new LLLocalSingleton(
-				))
+								getRocksDBColumn(db, cfh),
+								this::getSnapshotLambda,
+								LLLocalKeyValueDatabase.this.name,
+								name,
+								ColumnUtils.toString(singletonListColumnName),
+								dbWScheduler, dbRScheduler, defaultValue
+						);
+					} finally {
+						closeLock.unlockRead(closeReadLock);
+					}
+				})
 				.onErrorMap(cause -> new IOException("Failed to read " + Arrays.toString(name), cause))
 				.subscribeOn(dbRScheduler);
 	}
@@ -969,28 +1037,45 @@ public class LLLocalKeyValueDatabase implements LLKeyValueDatabase {
 	@Override
 	public Mono<LLLocalDictionary> getDictionary(byte[] columnName, UpdateMode updateMode) {
 		return Mono
-				.fromCallable(() -> new LLLocalDictionary(
+				.fromCallable(() -> {
-						allocator,
+					var closeReadLock = closeLock.readLock();
-						getRocksDBColumn(db, getCfh(columnName)),
+					try {
-						name,
+						ensureOpen();
-						ColumnUtils.toString(columnName),
+						var cfh = getCfh(columnName);
-						dbWScheduler,
+						ensureOwned(cfh);
-						dbRScheduler,
+						return new LLLocalDictionary(
-						(snapshot) -> snapshotsHandles.get(snapshot.getSequenceNumber()),
+								allocator,
-						updateMode,
+								getRocksDBColumn(db, cfh),
-						databaseOptions
+								name,
-				))
+								ColumnUtils.toString(columnName),
+								dbWScheduler,
+								dbRScheduler,
+								this::getSnapshotLambda,
+								updateMode,
+								databaseOptions
+						);
+					} finally {
+						closeLock.unlockRead(closeReadLock);
+					}
+				})
 				.subscribeOn(dbRScheduler);
 	}
 
 	public RocksDBColumn getRocksDBColumn(byte[] columnName) {
-		ColumnFamilyHandle cfh;
+		var closeReadLock = closeLock.readLock();
 		try {
-			cfh = getCfh(columnName);
+			ensureOpen();
-		} catch (RocksDBException e) {
+			ColumnFamilyHandle cfh;
-			throw new UnsupportedOperationException("Column family doesn't exist: " + Arrays.toString(columnName), e);
+			try {
+				cfh = getCfh(columnName);
+				ensureOwned(cfh);
+			} catch (RocksDBException e) {
+				throw new UnsupportedOperationException("Column family doesn't exist: " + Arrays.toString(columnName), e);
+			}
+			return getRocksDBColumn(db, cfh);
+		} finally {
+			closeLock.unlockRead(closeReadLock);
 		}
-		return getRocksDBColumn(db, cfh);
 	}
 
 	private RocksDBColumn getRocksDBColumn(RocksDB db, ColumnFamilyHandle cfh) {
@@ -1031,15 +1116,29 @@ public class LLLocalKeyValueDatabase implements LLKeyValueDatabase {
 
 	@Override
 	public Mono<Long> getProperty(String propertyName) {
-		return Mono.fromCallable(() -> db.getAggregatedLongProperty(propertyName))
+		return Mono.fromCallable(() -> {
-				.onErrorMap(cause -> new IOException("Failed to read " + propertyName, cause))
+			var closeReadLock = closeLock.readLock();
-				.subscribeOn(dbRScheduler);
+			try {
+				ensureOpen();
+				return db.getAggregatedLongProperty(propertyName);
+			} finally {
+				closeLock.unlockRead(closeReadLock);
+			}
+		}).onErrorMap(cause -> new IOException("Failed to read " + propertyName, cause)).subscribeOn(dbRScheduler);
 	}
 
 	public Flux<Path> getSSTS() {
 		var paths = convertPaths(dbPath.toAbsolutePath().getParent(), dbPath.getFileName(), databaseOptions.volumes());
 		return Mono
-				.fromCallable(() -> db.getLiveFiles())
+				.fromCallable(() -> {
+					var closeReadLock = closeLock.readLock();
+					try {
+						ensureOpen();
+						return db.getLiveFiles();
+					} finally {
+						closeLock.unlockRead(closeReadLock);
+					}
+				})
 				.flatMapIterable(liveFiles -> liveFiles.files)
 				.filter(file -> file.endsWith(".sst"))
 				.map(file -> file.substring(1))
@@ -1064,6 +1163,7 @@ public class LLLocalKeyValueDatabase implements LLKeyValueDatabase {
 		return sstsFlux
 				.map(path -> path.toAbsolutePath().toString())
 				.flatMap(sst -> Mono.fromCallable(() -> {
+					var closeReadLock = closeLock.readLock();
 					try (var opts = new IngestExternalFileOptions()) {
 						try {
 							logger.info("Ingesting SST \"{}\"...", sst);
@@ -1072,6 +1172,8 @@ public class LLLocalKeyValueDatabase implements LLKeyValueDatabase {
 						} catch (RocksDBException e) {
 							logger.error("Can't ingest SST \"{}\"", sst, e);
 						}
+					} finally {
+						closeLock.unlockRead(closeReadLock);
 					}
 					return null;
 				}).subscribeOn(Schedulers.boundedElastic()))
@@ -1082,16 +1184,23 @@ public class LLLocalKeyValueDatabase implements LLKeyValueDatabase {
 	public Mono<MemoryStats> getMemoryStats() {
 		return Mono
 				.fromCallable(() -> {
-					if (!closed) {
+					if (closed) return null;
-						return new MemoryStats(db.getAggregatedLongProperty("rocksdb.estimate-table-readers-mem"),
+					var closeReadLock = closeLock.readLock();
-								db.getAggregatedLongProperty("rocksdb.size-all-mem-tables"),
+					try {
-								db.getAggregatedLongProperty("rocksdb.cur-size-all-mem-tables"),
+						if (!closed) {
-								db.getAggregatedLongProperty("rocksdb.estimate-num-keys"),
+							ensureOpen();
-								db.getAggregatedLongProperty("rocksdb.block-cache-usage") / this.handles.size(),
+							return new MemoryStats(db.getAggregatedLongProperty("rocksdb.estimate-table-readers-mem"),
-								db.getAggregatedLongProperty("rocksdb.block-cache-pinned-usage") / this.handles.size()
+									db.getAggregatedLongProperty("rocksdb.size-all-mem-tables"),
-						);
+									db.getAggregatedLongProperty("rocksdb.cur-size-all-mem-tables"),
-					} else {
+									db.getAggregatedLongProperty("rocksdb.estimate-num-keys"),
-						return null;
+									db.getAggregatedLongProperty("rocksdb.block-cache-usage") / this.handles.size(),
+									db.getAggregatedLongProperty("rocksdb.block-cache-pinned-usage") / this.handles.size()
+							);
+						} else {
+							return null;
+						}
+					} finally {
+						closeLock.unlockRead(closeReadLock);
 					}
 				})
 				.onErrorMap(cause -> new IOException("Failed to read memory stats", cause))
@@ -1102,18 +1211,25 @@ public class LLLocalKeyValueDatabase implements LLKeyValueDatabase {
 	public Mono<String> getRocksDBStats() {
 		return Mono
 				.fromCallable(() -> {
-					if (!closed) {
+					if (closed) return null;
-						StringBuilder aggregatedStats = new StringBuilder();
+					var closeReadLock = closeLock.readLock();
-						for (var entry : this.handles.entrySet()) {
+					try {
-							aggregatedStats
+						if (!closed) {
-									.append(entry.getKey().name())
+							ensureOpen();
-									.append("\n")
+							StringBuilder aggregatedStats = new StringBuilder();
-									.append(db.getProperty(entry.getValue(), "rocksdb.stats"))
+							for (var entry : this.handles.entrySet()) {
-									.append("\n");
+								aggregatedStats
+										.append(entry.getKey().name())
+										.append("\n")
+										.append(db.getProperty(entry.getValue(), "rocksdb.stats"))
+										.append("\n");
+							}
+							return aggregatedStats.toString();
+						} else {
+							return null;
 						}
-						return aggregatedStats.toString();
+					} finally {
-					} else {
+						closeLock.unlockRead(closeReadLock);
-						return null;
 					}
 				})
 				.onErrorMap(cause -> new IOException("Failed to read stats", cause))
@@ -1126,10 +1242,14 @@ public class LLLocalKeyValueDatabase implements LLKeyValueDatabase {
 				.fromIterable(handles.entrySet())
 				.flatMapSequential(handle -> Mono
 						.fromCallable(() -> {
-							if (!closed) {
+							if (closed) return null;
+							var closeReadLock = closeLock.readLock();
+							try {
+								if (closed) return null;
+								ensureOpen();
 								return db.getPropertiesOfAllTables(handle.getValue());
-							} else {
+							} finally {
-								return null;
+								closeLock.unlockRead(closeReadLock);
 							}
 						})
 						.subscribeOn(dbRScheduler)
@@ -1143,7 +1263,13 @@ public class LLLocalKeyValueDatabase implements LLKeyValueDatabase {
 	public Mono<Void> verifyChecksum() {
 		return Mono
 				.<Void>fromCallable(() -> {
-					db.verifyChecksum();
+					var closeReadLock = closeLock.readLock();
+					try {
+						ensureOpen();
+						db.verifyChecksum();
+					} finally {
+						closeLock.unlockRead(closeReadLock);
+					}
 					return null;
 				})
 				.onErrorMap(cause -> new IOException("Failed to verify checksum of database \""
@@ -1181,22 +1307,20 @@ public class LLLocalKeyValueDatabase implements LLKeyValueDatabase {
 
 	@Override
 	public Mono<LLSnapshot> takeSnapshot() {
-		return Mono
+		return Mono.fromCallable(() -> {
-				.fromCallable(() -> snapshotTime.recordCallable(() -> {
+			var closeReadLock = closeLock.readLock();
-					var closeReadLock = closeLock.readLock();
+			try {
-					try {
+				ensureOpen();
-						if (closed) {
+				return snapshotTime.recordCallable(() -> {
-							throw new IllegalStateException("Database closed");
+					var snapshot = db.getSnapshot();
-						}
+					long currentSnapshotSequenceNumber = nextSnapshotNumbers.getAndIncrement();
-						var snapshot = db.getSnapshot();
+					this.snapshotsHandles.put(currentSnapshotSequenceNumber, snapshot);
-						long currentSnapshotSequenceNumber = nextSnapshotNumbers.getAndIncrement();
+					return new LLSnapshot(currentSnapshotSequenceNumber);
-						this.snapshotsHandles.put(currentSnapshotSequenceNumber, snapshot);
+				});
-						return new LLSnapshot(currentSnapshotSequenceNumber);
+			} finally {
-					} finally {
+				closeLock.unlockRead(closeReadLock);
-						closeLock.unlockRead(closeReadLock);
+			}
-					}
+		}).subscribeOn(dbRScheduler);
-				}))
-				.subscribeOn(dbRScheduler);
 	}
 
 	@Override
@@ -1205,9 +1329,6 @@ public class LLLocalKeyValueDatabase implements LLKeyValueDatabase {
 				.<Void>fromCallable(() -> {
 					var closeReadLock = closeLock.readLock();
 					try {
-						if (closed) {
-							throw new IllegalStateException("Database closed");
-						}
 						Snapshot dbSnapshot = this.snapshotsHandles.remove(snapshot.getSequenceNumber());
 						if (dbSnapshot == null) {
 							throw new IOException("Snapshot " + snapshot.getSequenceNumber() + " not found!");

src/main/java/it/cavallium/dbengine/database/disk/RocksDBUtils.java

@@ -7,6 +7,7 @@ import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.List;
 import org.apache.logging.log4j.Logger;
+import org.jetbrains.annotations.Nullable;
 import org.rocksdb.ColumnFamilyHandle;
 import org.rocksdb.CompactionJobInfo;
 import org.rocksdb.CompactionOptions;
@@ -94,4 +95,18 @@ public class RocksDBUtils {
 			}
 		}
 	}
+
+	public static void ensureOpen(RocksDB db, @Nullable ColumnFamilyHandle cfh) {
+		if (Schedulers.isInNonBlockingThread()) {
+			throw new UnsupportedOperationException("Called in a nonblocking thread");
+		}
+		ensureOwned(db);
+		ensureOwned(cfh);
+	}
+
+	public static void ensureOwned(@Nullable org.rocksdb.RocksObject rocksObject) {
+		if (rocksObject != null && !rocksObject.isOwningHandle()) {
+			throw new IllegalStateException("Not owning handle");
+		}
+	}
 }