Skip to content

Commit c378136

Browse files
BarbatosBarbatos
Barbatos
authored and
Barbatos
committed
fix(chainbase): thread safety and resource leak fixes
- Chainbase.head: declare as volatile for cross-thread visibility - Chainbase.prefixQuery(): snapshot head into local variable at entry and pass to prefixQueryRoot()/prefixQuerySnapshot() to guarantee consistency within a single call - TronDatabase.close(): split writeOptions.close() and dbSource.closeDB() into separate try-catch blocks so one failure does not skip the other - CheckPointV2Store.close(): close own writeOptions in try-catch then call super.close() for full parent cleanup chain - LevelDB.close() / RocksDB.close(): same split try-catch pattern - SnapshotManager.checkV2(): wrap getCheckpointDB() in try-with-resources so close() runs even if recover() throws
1 parent bd763bc commit c378136

11 files changed

Lines changed: 246 additions & 30 deletions

File tree

‎chainbase/src/main/java/org/tron/core/db/TronDatabase.java‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -78,12 +78,15 @@ public void close() {
7878
logger.info("******** Begin to close {}. ********", getName());
7979
try {
8080
writeOptions.close();
81+
} catch (Exception e) {
82+
logger.warn("Failed to close writeOptions in {}.", getName(), e);
83+
}
84+
try {
8185
dbSource.closeDB();
8286
} catch (Exception e) {
83-
logger.warn("Failed to close {}.", getName(), e);
84-
} finally {
85-
logger.info("******** End to close {}. ********", getName());
87+
logger.warn("Failed to close dbSource in {}.", getName(), e);
8688
}
89+
logger.info("******** End to close {}. ********", getName());
8790
}
8891

8992
public abstract void put(byte[] key, T item);

‎chainbase/src/main/java/org/tron/core/db2/common/LevelDB.java‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,11 +4,13 @@
44
import java.util.HashMap;
55
import java.util.Map;
66
import lombok.Getter;
7+
import lombok.extern.slf4j.Slf4j;
78
import org.tron.common.parameter.CommonParameter;
89
import org.tron.common.storage.WriteOptionsWrapper;
910
import org.tron.common.storage.leveldb.LevelDbDataSourceImpl;
1011
import org.tron.core.db.common.iterator.DBIterator;
1112

13+
@Slf4j(topic = "DB")
1214
public class LevelDB implements DB<byte[], byte[]>, Flusher {
1315

1416
@Getter
@@ -65,8 +67,16 @@ public void flush(Map<WrappedByteArray, WrappedByteArray> batch) {
6567

6668
@Override
6769
public void close() {
68-
this.writeOptions.close();
69-
db.closeDB();
70+
try {
71+
this.writeOptions.close();
72+
} catch (Exception e) {
73+
logger.warn("Failed to close writeOptions.", e);
74+
}
75+
try {
76+
db.closeDB();
77+
} catch (Exception e) {
78+
logger.warn("Failed to close db.", e);
79+
}
7080
}
7181

7282
@Override

‎chainbase/src/main/java/org/tron/core/db2/common/RocksDB.java‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,11 +4,13 @@
44
import java.util.HashMap;
55
import java.util.Map;
66
import lombok.Getter;
7+
import lombok.extern.slf4j.Slf4j;
78
import org.tron.common.parameter.CommonParameter;
89
import org.tron.common.storage.WriteOptionsWrapper;
910
import org.tron.common.storage.rocksdb.RocksDbDataSourceImpl;
1011
import org.tron.core.db.common.iterator.DBIterator;
1112

13+
@Slf4j(topic = "DB")
1214
public class RocksDB implements DB<byte[], byte[]>, Flusher {
1315

1416
@Getter
@@ -66,8 +68,16 @@ public void flush(Map<WrappedByteArray, WrappedByteArray> batch) {
6668

6769
@Override
6870
public void close() {
69-
writeOptions.close();
70-
db.closeDB();
71+
try {
72+
writeOptions.close();
73+
} catch (Exception e) {
74+
logger.warn("Failed to close writeOptions.", e);
75+
}
76+
try {
77+
db.closeDB();
78+
} catch (Exception e) {
79+
logger.warn("Failed to close db.", e);
80+
}
7181
}
7282

7383
@Override

‎chainbase/src/main/java/org/tron/core/db2/core/Chainbase.java‎

Lines changed: 15 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ public enum Cursor {
3535
//true:fullnode, false:soliditynode
3636
private ThreadLocal<Cursor> cursor = new ThreadLocal<>();
3737
private ThreadLocal<Long> offset = new ThreadLocal<>();
38-
private Snapshot head;
38+
private volatile Snapshot head;
3939

4040
public Chainbase(Snapshot head) {
4141
this.head = head;
@@ -349,28 +349,32 @@ private Map<byte[], byte[]> getNext(Snapshot head, byte[] key, long limit) {
349349
.collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
350350
}
351351

352+
// Snapshot head once: prefixQueryRoot and prefixQuerySnapshot both need the same
353+
// head reference. Other read methods (get, getUnchecked, has) only read head once
354+
// via head(), so volatile alone is sufficient for them.
352355
public Map<WrappedByteArray, byte[]> prefixQuery(byte[] key) {
353-
Map<WrappedByteArray, byte[]> result = prefixQueryRoot(key);
354-
Map<WrappedByteArray, byte[]> snapshot = prefixQuerySnapshot(key);
356+
Snapshot localHead = head;
357+
Map<WrappedByteArray, byte[]> result = prefixQueryRoot(localHead, key);
358+
Map<WrappedByteArray, byte[]> snapshot = prefixQuerySnapshot(localHead, key);
355359
result.putAll(snapshot);
356360
result.entrySet().removeIf(e -> e.getValue() == null);
357361
return result;
358362
}
359363

360-
private Map<WrappedByteArray, byte[]> prefixQueryRoot(byte[] key) {
364+
private Map<WrappedByteArray, byte[]> prefixQueryRoot(Snapshot localHead, byte[] key) {
361365
Map<WrappedByteArray, byte[]> result = new HashMap<>();
362-
if (((SnapshotRoot) head.getRoot()).db.getClass() == LevelDB.class) {
363-
result = ((LevelDB) ((SnapshotRoot) head.getRoot()).db).getDb().prefixQuery(key);
364-
} else if (((SnapshotRoot) head.getRoot()).db.getClass() == RocksDB.class) {
365-
result = ((RocksDB) ((SnapshotRoot) head.getRoot()).db).getDb().prefixQuery(key);
366+
if (((SnapshotRoot) localHead.getRoot()).db.getClass() == LevelDB.class) {
367+
result = ((LevelDB) ((SnapshotRoot) localHead.getRoot()).db).getDb().prefixQuery(key);
368+
} else if (((SnapshotRoot) localHead.getRoot()).db.getClass() == RocksDB.class) {
369+
result = ((RocksDB) ((SnapshotRoot) localHead.getRoot()).db).getDb().prefixQuery(key);
366370
}
367371
return result;
368372
}
369373

370-
private Map<WrappedByteArray, byte[]> prefixQuerySnapshot(byte[] key) {
374+
private Map<WrappedByteArray, byte[]> prefixQuerySnapshot(Snapshot localHead, byte[] key) {
371375
Map<WrappedByteArray, byte[]> result = new HashMap<>();
372-
Snapshot snapshot = head();
373-
if (!snapshot.equals(head.getRoot())) {
376+
Snapshot snapshot = localHead;
377+
if (!snapshot.equals(localHead.getRoot())) {
374378
Map<WrappedByteArray, WrappedByteArray> all = new HashMap<>();
375379
((SnapshotImpl) snapshot).collect(all, key);
376380
all.forEach((k, v) -> result.put(k, v.getBytes()));

‎chainbase/src/main/java/org/tron/core/db2/core/SnapshotManager.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -511,9 +511,9 @@ private void checkV2() {
511511
if (latestTimestamp - timestamp > ONE_MINUTE_MILLS*2) {
512512
continue;
513513
}
514-
TronDatabase<byte[]> checkPointV2Store = getCheckpointDB(cp);
515-
recover(checkPointV2Store);
516-
checkPointV2Store.close();
514+
try (TronDatabase<byte[]> checkPointV2Store = getCheckpointDB(cp)) {
515+
recover(checkPointV2Store);
516+
}
517517
}
518518
logger.info("checkpoint v2 recover success");
519519
unChecked = false;

‎chainbase/src/main/java/org/tron/core/store/CheckPointV2Store.java‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -63,19 +63,19 @@ public void updateByBatch(Map<byte[], byte[]> rows) {
6363
}
6464

6565
/**
66-
* close the database.
66+
* Closes the database and releases all resources.
67+
* Note: this class and the parent TronDatabase each declare their own
68+
* writeOptions field (Java fields are not polymorphic). Both must be
69+
* closed: this.writeOptions here, and the parent's via super.close().
6770
*/
6871
@Override
6972
public void close() {
70-
logger.debug("******** Begin to close {}. ********", getName());
7173
try {
7274
writeOptions.close();
73-
dbSource.closeDB();
7475
} catch (Exception e) {
75-
logger.warn("Failed to close {}.", getName(), e);
76-
} finally {
77-
logger.debug("******** End to close {}. ********", getName());
76+
logger.warn("Failed to close writeOptions in {}.", getName(), e);
7877
}
78+
super.close();
7979
}
8080

8181
}

‎framework/src/main/java/org/tron/core/net/service/effective/ResilienceService.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -246,6 +246,10 @@ public Object next() {
246246
for (int weight : weights.values()) {
247247
totalWeight += weight;
248248
}
249+
if (totalWeight <= 0) {
250+
throw new IllegalStateException(
251+
"Total weight must be positive, got: " + totalWeight);
252+
}
249253
int randomNum = ThreadLocalRandom.current().nextInt(totalWeight);
250254
for (Object key : weights.keySet()) {
251255
randomNum -= weights.get(key);

‎framework/src/test/java/org/tron/core/db/DefensiveChecksTest.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ public class DefensiveChecksTest extends BaseTest {
2828
// -------------------------------------------------------------------------
2929

3030
@Test
31-
public void testEnergyProcessor_calculateGlobalEnergyLimit_zeroWeight_throwsWhenNewRewardDisabled() {
31+
public void testEnergyProcessorZeroWeightThrowsWhenNewRewardDisabled() {
3232
// Arrange: new-reward feature off, energy weight = 0, account with enough frozen balance
3333
dbManager.getDynamicPropertiesStore().saveUnfreezeDelayDays(0L);
3434
dbManager.getDynamicPropertiesStore().saveAllowNewReward(0L);
@@ -50,7 +50,7 @@ public void testEnergyProcessor_calculateGlobalEnergyLimit_zeroWeight_throwsWhen
5050
}
5151

5252
@Test
53-
public void testEnergyProcessor_calculateGlobalEnergyLimit_zeroWeight_returnsZeroWhenNewRewardEnabled() {
53+
public void testEnergyProcessorZeroWeightReturnsZeroWhenNewRewardEnabled() {
5454
// When allowNewReward is on, totalEnergyWeight == 0 should return 0 (not throw)
5555
dbManager.getDynamicPropertiesStore().saveUnfreezeDelayDays(0L);
5656
dbManager.getDynamicPropertiesStore().saveAllowNewReward(1L);

‎framework/src/test/java/org/tron/core/db/ResourceCloseTest.java‎

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,86 @@
1+
package org.tron.core.db;
2+
3+
import static org.mockito.Mockito.doThrow;
4+
import static org.mockito.Mockito.mock;
5+
import static org.mockito.Mockito.spy;
6+
import static org.mockito.Mockito.verify;
7+
8+
import java.lang.reflect.Field;
9+
import org.junit.AfterClass;
10+
import org.junit.BeforeClass;
11+
import org.junit.ClassRule;
12+
import org.junit.Test;
13+
import org.junit.rules.TemporaryFolder;
14+
import org.tron.common.TestConstants;
15+
import org.tron.common.storage.WriteOptionsWrapper;
16+
import org.tron.core.config.args.Args;
17+
import org.tron.core.db.common.DbSourceInter;
18+
import org.tron.core.store.CheckPointV2Store;
19+
20+
/**
21+
* Verifies that close() methods properly release all resources even when
22+
* one resource's close() throws an exception.
23+
*/
24+
public class ResourceCloseTest {
25+
26+
@ClassRule
27+
public static final TemporaryFolder temporaryFolder = new TemporaryFolder();
28+
29+
static {
30+
org.rocksdb.RocksDB.loadLibrary();
31+
}
32+
33+
@BeforeClass
34+
public static void init() throws Exception {
35+
Args.setParam(
36+
new String[]{"-d", temporaryFolder.newFolder().toString()},
37+
TestConstants.TEST_CONF);
38+
}
39+
40+
@AfterClass
41+
public static void destroy() {
42+
Args.clearParam();
43+
}
44+
45+
@Test
46+
public void testTronDatabase_closeDbSource_whenWriteOptionsThrows() throws Exception {
47+
CheckPointV2Store store = new CheckPointV2Store("test-close-safety");
48+
49+
// Get parent class fields via reflection
50+
Field writeOptionsField = TronDatabase.class.getDeclaredField("writeOptions");
51+
writeOptionsField.setAccessible(true);
52+
WriteOptionsWrapper originalWriteOptions =
53+
(WriteOptionsWrapper) writeOptionsField.get(store);
54+
55+
Field dbSourceField = TronDatabase.class.getDeclaredField("dbSource");
56+
dbSourceField.setAccessible(true);
57+
DbSourceInter<?> originalDbSource =
58+
(DbSourceInter<?>) dbSourceField.get(store);
59+
60+
// Replace with spies
61+
WriteOptionsWrapper spyWriteOptions = spy(originalWriteOptions);
62+
doThrow(new RuntimeException("writeOptions close failed"))
63+
.when(spyWriteOptions).close();
64+
writeOptionsField.set(store, spyWriteOptions);
65+
66+
DbSourceInter<?> spyDbSource = spy(originalDbSource);
67+
dbSourceField.set(store, spyDbSource);
68+
69+
// Also spy the child's writeOptions
70+
Field childWriteOptionsField = CheckPointV2Store.class.getDeclaredField("writeOptions");
71+
childWriteOptionsField.setAccessible(true);
72+
WriteOptionsWrapper childOriginal =
73+
(WriteOptionsWrapper) childWriteOptionsField.get(store);
74+
WriteOptionsWrapper spyChildWriteOptions = spy(childOriginal);
75+
doThrow(new RuntimeException("child writeOptions close failed"))
76+
.when(spyChildWriteOptions).close();
77+
childWriteOptionsField.set(store, spyChildWriteOptions);
78+
79+
// close() should not throw, and dbSource should still be closed
80+
store.close();
81+
82+
verify(spyChildWriteOptions).close();
83+
verify(spyWriteOptions).close();
84+
verify(spyDbSource).closeDB();
85+
}
86+
}

‎framework/src/test/java/org/tron/core/db/TransactionTraceTest.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,8 +45,8 @@
4545
import org.tron.protos.Protocol.Transaction.Contract;
4646
import org.tron.protos.Protocol.Transaction.Contract.ContractType;
4747
import org.tron.protos.Protocol.Transaction.raw;
48-
import org.tron.protos.contract.SmartContractOuterClass.CreateSmartContract;
4948
import org.tron.protos.contract.Common.ResourceCode;
49+
import org.tron.protos.contract.SmartContractOuterClass.CreateSmartContract;
5050
import org.tron.protos.contract.SmartContractOuterClass.SmartContract;
5151
import org.tron.protos.contract.SmartContractOuterClass.TriggerSmartContract;
5252

0 commit comments

Comments
 (0)