Skip to content

Commit 3272196

Browse files
EAOE3EAOE3
committed
Refactor closed flag in MerkleTree to use AtomicBoolean for thread-safe access and add error handling for closed state
1 parent de50535 commit 3272196

File tree

1 file changed

+28
-178
lines changed

1 file changed

+28
-178
lines changed

‎src/main/java/io/pwrlabs/database/rocksdb/MerkleTree.java‎

Lines changed: 28 additions & 178 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
import java.nio.ByteBuffer;
1414
import java.util.*;
1515
import java.util.concurrent.ConcurrentHashMap;
16+
import java.util.concurrent.atomic.AtomicBoolean;
1617
import java.util.concurrent.locks.ReadWriteLock;
1718
import java.util.concurrent.locks.ReentrantReadWriteLock;
1819

@@ -71,181 +72,10 @@ public class MerkleTree {
7172
private int depth = 0;
7273
private byte[] rootHash = null;
7374

74-
private boolean closed = false;
75+
private AtomicBoolean closed = new AtomicBoolean(false);
7576
//endregion
7677

7778
//region ===================== Constructors =====================
78-
// @SneakyThrows
79-
// public MerkleTree(String treeName) {
80-
// errorIf(openTrees.containsKey(treeName), "There is already open instance of this tree. 2 open instances of the same tree are not allowed at the same time");
81-
// this.treeName = treeName;
82-
//
83-
// Options options = new Options().setCreateIfMissing(true)
84-
// .setMaxTotalWalSize(45 * 1024 * 1024L) // Good: Limits total WAL to 512MB
85-
// .setWalSizeLimitMB(15) // Good: Standard size per WAL file
86-
// .setWalTtlSeconds(24 * 60 * 60) // Good: Cleans up old WAL files after 24h
87-
// .setInfoLogLevel(InfoLogLevel.FATAL_LEVEL) // Good: Minimizes logging overhead
88-
// .setDbLogDir("") // Good: Disables separate log directory
89-
// .setLogFileTimeToRoll(0); // Good: Immediate roll when size limit reached
90-
//
91-
// // Add these additional safety options
92-
// options.setAllowMmapReads(false) // Disable memory mapping
93-
// .setAllowMmapWrites(false)
94-
// .setMaxOpenFiles(1000)
95-
// .setMaxFileOpeningThreads(10)
96-
// .setIncreaseParallelism(1); // Single-threaded mode is safer
97-
//
98-
// options.setParanoidChecks(true) // Enable paranoid checks for corruption
99-
// .setUseDirectReads(true) // Direct I/O for reads
100-
// .setUseDirectIoForFlushAndCompaction(true) // Direct I/O for writes
101-
// .setEnableWriteThreadAdaptiveYield(true)
102-
// .setAllowConcurrentMemtableWrite(true);
103-
//
104-
// File directory = new File("merkleTree/");
105-
// if(!directory.exists()) directory.mkdirs();
106-
//
107-
// this.db = RocksDB.open(options, "merkleTree/" + treeName);
108-
//
109-
// final ColumnFamilyDescriptor db1Descriptor =
110-
// new ColumnFamilyDescriptor(METADATA_DB_NAME.getBytes());
111-
// final ColumnFamilyDescriptor db2Descriptor =
112-
// new ColumnFamilyDescriptor(NODES_DB_NAME.getBytes());
113-
//
114-
// metaDataHandle = db.createColumnFamily(db1Descriptor);
115-
// nodesHandle = db.createColumnFamily(db2Descriptor);
116-
//
117-
// loadMetaData();
118-
//
119-
// //Shutdown hook
120-
// Runtime.getRuntime().addShutdownHook(new Thread(() -> {
121-
// try {
122-
// close();
123-
// } catch (RocksDBException e) {
124-
// e.printStackTrace();
125-
// }
126-
// })
127-
// );
128-
//
129-
// openTrees.put(treeName, this);
130-
// }
131-
132-
// public MerkleTree(String treeName) throws RocksDBException {
133-
// RocksDB.loadLibrary();
134-
// this.treeName = treeName;
135-
// errorIf(openTrees.containsKey(treeName), "There is already open instance of this tree. 2 open instances of the same tree are not allowed at the same time");
136-
//
137-
// DBOptions dbOptions = new DBOptions().setCreateIfMissing(true)
138-
// .setMaxTotalWalSize(45 * 1024 * 1024L) // Good: Limits total WAL to 512MB
139-
// .setWalSizeLimitMB(15) // Good: Standard size per WAL file
140-
// .setWalTtlSeconds(24 * 60 * 60) // Good: Cleans up old WAL files after 24h
141-
// .setInfoLogLevel(InfoLogLevel.FATAL_LEVEL) // Good: Minimizes logging overhead
142-
// .setDbLogDir("") // Good: Disables separate log directory
143-
// .setLogFileTimeToRoll(0); // Good: Immediate roll when size limit reached
144-
//
145-
// // Add these additional safety options
146-
// dbOptions.setAllowMmapReads(false) // Disable memory mapping
147-
// .setAllowMmapWrites(false)
148-
// .setMaxOpenFiles(1000)
149-
// .setMaxFileOpeningThreads(10)
150-
// .setIncreaseParallelism(1); // Single-threaded mode is safer
151-
//
152-
// dbOptions.setParanoidChecks(true) // Enable paranoid checks for corruption
153-
// .setUseDirectReads(true) // Direct I/O for reads
154-
// .setUseDirectIoForFlushAndCompaction(true) // Direct I/O for writes
155-
// .setEnableWriteThreadAdaptiveYield(true)
156-
// .setAllowConcurrentMemtableWrite(true);
157-
// // set your other options here...
158-
//
159-
// String dbPath = "merkleTree/" + treeName;
160-
//
161-
//
162-
// Options options = new Options().setCreateIfMissing(true)
163-
// .setMaxTotalWalSize(45 * 1024 * 1024L) // Good: Limits total WAL to 512MB
164-
// .setWalSizeLimitMB(15) // Good: Standard size per WAL file
165-
// .setWalTtlSeconds(24 * 60 * 60) // Good: Cleans up old WAL files after 24h
166-
// .setInfoLogLevel(InfoLogLevel.FATAL_LEVEL) // Good: Minimizes logging overhead
167-
// .setDbLogDir("") // Good: Disables separate log directory
168-
// .setLogFileTimeToRoll(0); // Good: Immediate roll when size limit reached
169-
//
170-
// // Add these additional safety options
171-
// options.setAllowMmapReads(false) // Disable memory mapping
172-
// .setAllowMmapWrites(false)
173-
// .setMaxOpenFiles(1000)
174-
// .setMaxFileOpeningThreads(10)
175-
// .setIncreaseParallelism(1); // Single-threaded mode is safer
176-
//
177-
// options.setParanoidChecks(true) // Enable paranoid checks for corruption
178-
// .setUseDirectReads(true) // Direct I/O for reads
179-
// .setUseDirectIoForFlushAndCompaction(true) // Direct I/O for writes
180-
// .setEnableWriteThreadAdaptiveYield(true)
181-
// .setAllowConcurrentMemtableWrite(true);
182-
//
183-
// // 1) Figure out which column families already exist
184-
// List<byte[]> existingCFNames = RocksDB.listColumnFamilies(options, dbPath);
185-
//
186-
// List<ColumnFamilyDescriptor> cfDescriptors = new ArrayList<>();
187-
// if (existingCFNames.isEmpty()) {
188-
// // Means this is a brand new DB -- no column families yet
189-
// // We always need the default CF
190-
// cfDescriptors.add(new ColumnFamilyDescriptor(
191-
// RocksDB.DEFAULT_COLUMN_FAMILY,
192-
// new ColumnFamilyOptions())
193-
// );
194-
//
195-
// // Also create metaData CF
196-
// cfDescriptors.add(new ColumnFamilyDescriptor(
197-
// METADATA_DB_NAME.getBytes(),
198-
// new ColumnFamilyOptions())
199-
// );
200-
//
201-
// // Also create nodes CF
202-
// cfDescriptors.add(new ColumnFamilyDescriptor(
203-
// NODES_DB_NAME.getBytes(),
204-
// new ColumnFamilyOptions())
205-
// );
206-
// } else {
207-
// // We already have some (or all) CFs in the DB. We must open them *all*.
208-
// for (byte[] cfName : existingCFNames) {
209-
// cfDescriptors.add(
210-
// new ColumnFamilyDescriptor(cfName, new ColumnFamilyOptions())
211-
// );
212-
// }
213-
// }
214-
//
215-
// // 2) Open DB with all column family descriptors
216-
// List<ColumnFamilyHandle> cfHandles = new ArrayList<>();
217-
// this.db = RocksDB.open(dbOptions, dbPath, cfDescriptors, cfHandles);
218-
//
219-
// // 3) Figure out which handle corresponds to metaData, which to nodes
220-
// // They come back in the same order we put them in cfDescriptors.
221-
// for (int i = 0; i < cfDescriptors.size(); i++) {
222-
// String cfName = new String(cfDescriptors.get(i).getName());
223-
// if (cfName.equals(METADATA_DB_NAME)) {
224-
// metaDataHandle = cfHandles.get(i);
225-
// } else if (cfName.equals(NODES_DB_NAME)) {
226-
// nodesHandle = cfHandles.get(i);
227-
// } else if (cfName.equals("default")) {
228-
// // If you need the default CF handle, grab it here
229-
// }
230-
// }
231-
//
232-
// // If we found that we do NOT have metaDataHandle or nodesHandle yet (for example, an existing DB
233-
// // had only a default CF), you can create them here:
234-
// if (metaDataHandle == null) {
235-
// metaDataHandle = db.createColumnFamily(
236-
// new ColumnFamilyDescriptor(METADATA_DB_NAME.getBytes(), new ColumnFamilyOptions())
237-
// );
238-
// }
239-
// if (nodesHandle == null) {
240-
// nodesHandle = db.createColumnFamily(
241-
// new ColumnFamilyDescriptor(NODES_DB_NAME.getBytes(), new ColumnFamilyOptions())
242-
// );
243-
// }
244-
//
245-
// // 4) Proceed as normal (e.g. load metadata, etc.)
246-
// loadMetaData();
247-
// }
248-
24979
public MerkleTree(String treeName) throws RocksDBException {
25080
RocksDB.loadLibrary();
25181
this.treeName = treeName;
@@ -337,6 +167,7 @@ public MerkleTree(String treeName) throws RocksDBException {
337167
* Get the current root hash of the Merkle tree.
338168
*/
339169
public byte[] getRootHash() {
170+
errorIfClosed();
340171
lock.readLock().lock();
341172
try {
342173
if(rootHash == null) return null;
@@ -347,6 +178,7 @@ public byte[] getRootHash() {
347178
}
348179

349180
public int getNumLeaves() {
181+
errorIfClosed();
350182
lock.readLock().lock();
351183
try {
352184
return numLeaves;
@@ -356,6 +188,7 @@ public int getNumLeaves() {
356188
}
357189

358190
public int getDepth() {
191+
errorIfClosed();
359192
lock.readLock().lock();
360193
try {
361194
return depth;
@@ -370,6 +203,7 @@ public int getDepth() {
370203
* @throws RocksDBException If there's an error accessing RocksDB
371204
*/
372205
public HashSet<Node> getAllNodes() throws RocksDBException {
206+
errorIfClosed();
373207
lock.readLock().lock();
374208
try {
375209
HashSet<Node> allNodes = new HashSet<>();
@@ -406,6 +240,7 @@ public HashSet<Node> getAllNodes() throws RocksDBException {
406240
* @throws IllegalArgumentException If key is null
407241
*/
408242
public byte[] getData(byte[] key) {
243+
errorIfClosed();
409244
byte[] data = keyDataCache.get(new ByteArrayWrapper(key));
410245
if(data != null) return data;
411246

@@ -427,6 +262,8 @@ public byte[] getData(byte[] key) {
427262
* @throws IllegalArgumentException If key or data is null
428263
*/
429264
public void addOrUpdateData(byte[] key, byte[] data) throws RocksDBException {
265+
errorIfClosed();
266+
430267
if (key == null) {
431268
throw new IllegalArgumentException("Key cannot be null");
432269
}
@@ -464,6 +301,8 @@ public void addOrUpdateData(byte[] key, byte[] data) throws RocksDBException {
464301
}
465302

466303
public void revertUnsavedChanges() {
304+
errorIfClosed();
305+
467306
lock.writeLock().lock();
468307
try {
469308
nodesCache.clear();
@@ -479,6 +318,8 @@ public void revertUnsavedChanges() {
479318
}
480319

481320
public boolean containsKey(byte[] key) {
321+
errorIfClosed();
322+
482323
if (key == null) {
483324
throw new IllegalArgumentException("Key cannot be null");
484325
}
@@ -494,6 +335,7 @@ public boolean containsKey(byte[] key) {
494335
}
495336

496337
public List<byte[]> getAllKeys() {
338+
errorIfClosed();
497339
lock.readLock().lock();
498340
try {
499341
List<byte[]> keys = new ArrayList<>();
@@ -511,6 +353,7 @@ public List<byte[]> getAllKeys() {
511353
}
512354

513355
public List<byte[]> getAllData() {
356+
errorIfClosed();
514357
lock.readLock().lock();
515358
try {
516359
List<byte[]> data = new ArrayList<>();
@@ -531,6 +374,7 @@ public List<byte[]> getAllData() {
531374
* Flush all in-memory changes (nodes, metadata) to RocksDB.
532375
*/
533376
public void flushToDisk() throws RocksDBException {
377+
errorIfClosed();
534378
lock.writeLock().lock();
535379
try {
536380
try (WriteBatch batch = new WriteBatch()) {
@@ -584,7 +428,7 @@ public void flushToDisk() throws RocksDBException {
584428
public void close() throws RocksDBException {
585429
lock.writeLock().lock();
586430
try {
587-
if(closed) return;
431+
if(closed.get()) return;
588432
flushToDisk();
589433

590434
if (metaDataHandle != null) {
@@ -620,13 +464,15 @@ public void close() throws RocksDBException {
620464
}
621465

622466
openTrees.remove(treeName);
623-
closed = true;
467+
closed.set(true);
624468
} finally {
625469
lock.writeLock().unlock();
626470
}
627471
}
628472

629473
public MerkleTree clone(String newTreeName) throws RocksDBException {
474+
errorIfClosed();
475+
630476
if(newTreeName == null || newTreeName.isEmpty()) {
631477
throw new IllegalArgumentException("New tree name cannot be null or empty");
632478
}
@@ -670,10 +516,11 @@ public MerkleTree clone(String newTreeName) throws RocksDBException {
670516
* This is much faster than iterating through all entries and deleting them individually.
671517
*/
672518
public void clear() {
519+
errorIfClosed();
673520
lock.writeLock().lock();
674521
try {
675522
// First close the current DB
676-
if (!closed) {
523+
if (!closed.get()) {
677524
// Close all column family handles
678525
if (metaDataHandle != null) {
679526
try {
@@ -783,9 +630,6 @@ public void clear() {
783630
this.metaDataHandle = cfHandles.get(1);
784631
this.nodesHandle = cfHandles.get(2);
785632
this.keyDataHandle = cfHandles.get(3);
786-
787-
// Reset closed flag
788-
closed = false;
789633
} catch (RocksDBException e) {
790634
throw new RuntimeException("Failed to re-initialize RocksDB after clearing: " + e.getMessage(), e);
791635
}
@@ -1021,6 +865,12 @@ private void addNode(int level, Node node) throws RocksDBException {
1021865
}
1022866
}
1023867

868+
private void errorIfClosed() {
869+
if (closed.get()) {
870+
throw new IllegalStateException("MerkleTree is closed");
871+
}
872+
}
873+
1024874
//endregion
1025875

1026876
//region ===================== Nested Classes =====================

0 commit comments

Comments
 (0)