diff --git a/chainbase/src/main/java/org/tron/core/db/TronDatabase.java b/chainbase/src/main/java/org/tron/core/db/TronDatabase.java index 40762568c82..82f47d311e1 100644 --- a/chainbase/src/main/java/org/tron/core/db/TronDatabase.java +++ b/chainbase/src/main/java/org/tron/core/db/TronDatabase.java @@ -76,13 +76,23 @@ public void reset() { @Override public void close() { logger.info("******** Begin to close {}. ********", getName()); + try { + doClose(); + } finally { + logger.info("******** End to close {}. ********", getName()); + } + } + + protected void doClose() { try { writeOptions.close(); + } catch (Exception e) { + logger.warn("Failed to close writeOptions in {}.", getName(), e); + } + try { dbSource.closeDB(); } catch (Exception e) { - logger.warn("Failed to close {}.", getName(), e); - } finally { - logger.info("******** End to close {}. ********", getName()); + logger.warn("Failed to close dbSource in {}.", getName(), e); } } diff --git a/chainbase/src/main/java/org/tron/core/store/CheckPointV2Store.java b/chainbase/src/main/java/org/tron/core/store/CheckPointV2Store.java index f027bd02664..b886192f262 100644 --- a/chainbase/src/main/java/org/tron/core/store/CheckPointV2Store.java +++ b/chainbase/src/main/java/org/tron/core/store/CheckPointV2Store.java @@ -62,17 +62,16 @@ public void updateByBatch(Map rows) { this.dbSource.updateByBatch(rows, writeOptions); } - /** - * close the database. - */ @Override public void close() { logger.debug("******** Begin to close {}. ********", getName()); try { writeOptions.close(); - dbSource.closeDB(); } catch (Exception e) { - logger.warn("Failed to close {}.", getName(), e); + logger.warn("Failed to close writeOptions in {}.", getName(), e); + } + try { + doClose(); } finally { logger.debug("******** End to close {}. ********", getName()); } diff --git a/framework/src/test/java/org/tron/core/db/CheckPointV2StoreTest.java b/framework/src/test/java/org/tron/core/db/CheckPointV2StoreTest.java new file mode 100644 index 00000000000..4f90cf46b1a --- /dev/null +++ b/framework/src/test/java/org/tron/core/db/CheckPointV2StoreTest.java @@ -0,0 +1,129 @@ +package org.tron.core.db; + +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.verify; + +import java.io.IOException; +import java.lang.reflect.Field; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.ClassRule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; +import org.rocksdb.RocksDB; +import org.tron.common.TestConstants; +import org.tron.common.storage.WriteOptionsWrapper; +import org.tron.core.config.args.Args; +import org.tron.core.db.common.DbSourceInter; +import org.tron.core.store.CheckPointV2Store; + +public class CheckPointV2StoreTest { + + @ClassRule + public static final TemporaryFolder temporaryFolder = new TemporaryFolder(); + + static { + RocksDB.loadLibrary(); + } + + @BeforeClass + public static void initArgs() throws IOException { + Args.setParam( + new String[]{"-d", temporaryFolder.newFolder().toString()}, + TestConstants.TEST_CONF + ); + } + + @AfterClass + public static void destroy() { + Args.clearParam(); + } + + @Test + public void testCloseWithRealResources() { + CheckPointV2Store store = new CheckPointV2Store("test-close-real"); + // Exercises the real writeOptions.close() and dbSource.closeDB() code paths + store.close(); + } + + @Test + public void testCloseReleasesAllResources() throws Exception { + CheckPointV2Store store = new CheckPointV2Store("test-close"); + + // Replace dbSource with a mock so we can verify closeDB() + Field dbSourceField = TronDatabase.class.getDeclaredField("dbSource"); + dbSourceField.setAccessible(true); + DbSourceInter originalDbSource = (DbSourceInter) dbSourceField.get(store); + DbSourceInter mockDbSource = mock(DbSourceInter.class); + dbSourceField.set(store, mockDbSource); + + try { + store.close(); + + verify(mockDbSource).closeDB(); + } finally { + originalDbSource.closeDB(); + } + } + + @Test + public void testCloseWhenDbSourceThrows() throws Exception { + CheckPointV2Store store = new CheckPointV2Store("test-close-dbsource-throws"); + + Field dbSourceField = TronDatabase.class.getDeclaredField("dbSource"); + dbSourceField.setAccessible(true); + DbSourceInter originalDbSource = (DbSourceInter) dbSourceField.get(store); + DbSourceInter mockDbSource = mock(DbSourceInter.class); + doThrow(new RuntimeException("simulated dbSource failure")).when(mockDbSource).closeDB(); + dbSourceField.set(store, mockDbSource); + + try { + store.close(); + } finally { + originalDbSource.closeDB(); + } + } + + @Test + public void testCloseDbSourceWhenWriteOptionsThrows() throws Exception { + CheckPointV2Store store = new CheckPointV2Store("test-close-exception"); + + // Replace child writeOptions with a spy that throws on close + Field childWriteOptionsField = CheckPointV2Store.class.getDeclaredField("writeOptions"); + childWriteOptionsField.setAccessible(true); + WriteOptionsWrapper childWriteOptions = + (WriteOptionsWrapper) childWriteOptionsField.get(store); + WriteOptionsWrapper spyChildWriteOptions = spy(childWriteOptions); + doThrow(new RuntimeException("simulated writeOptions failure")) + .when(spyChildWriteOptions).close(); + childWriteOptionsField.set(store, spyChildWriteOptions); + + // Replace parent writeOptions with a spy that throws on close + Field parentWriteOptionsField = TronDatabase.class.getDeclaredField("writeOptions"); + parentWriteOptionsField.setAccessible(true); + WriteOptionsWrapper parentWriteOptions = + (WriteOptionsWrapper) parentWriteOptionsField.get(store); + WriteOptionsWrapper spyParentWriteOptions = spy(parentWriteOptions); + doThrow(new RuntimeException("simulated parent writeOptions failure")) + .when(spyParentWriteOptions).close(); + parentWriteOptionsField.set(store, spyParentWriteOptions); + + // Replace dbSource with a mock + Field dbSourceField = TronDatabase.class.getDeclaredField("dbSource"); + dbSourceField.setAccessible(true); + DbSourceInter originalDbSource = (DbSourceInter) dbSourceField.get(store); + DbSourceInter mockDbSource = mock(DbSourceInter.class); + dbSourceField.set(store, mockDbSource); + + try { + store.close(); + + // dbSource.closeDB() must be called even though both writeOptions threw + verify(mockDbSource).closeDB(); + } finally { + originalDbSource.closeDB(); + } + } +} diff --git a/framework/src/test/java/org/tron/core/db/TronDatabaseTest.java b/framework/src/test/java/org/tron/core/db/TronDatabaseTest.java index 52adf8e49fa..6bf479c287f 100644 --- a/framework/src/test/java/org/tron/core/db/TronDatabaseTest.java +++ b/framework/src/test/java/org/tron/core/db/TronDatabaseTest.java @@ -1,7 +1,13 @@ package org.tron.core.db; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.verify; + import com.google.protobuf.InvalidProtocolBufferException; import java.io.IOException; +import java.lang.reflect.Field; import org.junit.AfterClass; import org.junit.Assert; import org.junit.BeforeClass; @@ -12,7 +18,9 @@ import org.junit.rules.TemporaryFolder; import org.rocksdb.RocksDB; import org.tron.common.TestConstants; +import org.tron.common.storage.WriteOptionsWrapper; import org.tron.core.config.args.Args; +import org.tron.core.db.common.DbSourceInter; import org.tron.core.exception.BadItemException; import org.tron.core.exception.ItemNotFoundException; @@ -99,4 +107,48 @@ public void TestGetFromRoot() throws Assert.assertEquals(db.getFromRoot("test".getBytes()), "test"); } + + @Test + public void testDoCloseDbSourceCalledWhenWriteOptionsThrows() throws Exception { + TronDatabase db = new TronDatabase("test-do-close") { + + @Override + public void put(byte[] key, String item) { + } + + @Override + public void delete(byte[] key) { + } + + @Override + public String get(byte[] key) { + return null; + } + + @Override + public boolean has(byte[] key) { + return false; + } + }; + + Field writeOptionsField = TronDatabase.class.getDeclaredField("writeOptions"); + writeOptionsField.setAccessible(true); + WriteOptionsWrapper spyWriteOptions = spy((WriteOptionsWrapper) writeOptionsField.get(db)); + doThrow(new RuntimeException("simulated writeOptions failure")).when(spyWriteOptions).close(); + writeOptionsField.set(db, spyWriteOptions); + + Field dbSourceField = TronDatabase.class.getDeclaredField("dbSource"); + dbSourceField.setAccessible(true); + DbSourceInter originalDbSource = (DbSourceInter) dbSourceField.get(db); + DbSourceInter mockDbSource = mock(DbSourceInter.class); + dbSourceField.set(db, mockDbSource); + + try { + db.doClose(); + verify(spyWriteOptions).close(); + verify(mockDbSource).closeDB(); + } finally { + originalDbSource.closeDB(); + } + } }