fix(toolkit): keep db mv placement when lite merge replaces archive dbs

plugins/src/main/java/common/org/tron/plugins/DbLite.java

@@ -9,11 +9,15 @@
 import java.io.File;
 import java.io.FileNotFoundException;
 import java.io.IOException;
+import java.io.UncheckedIOException;
+import java.nio.file.FileVisitOption;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.util.Arrays;
+import java.util.Comparator;
 import java.util.List;
+import java.util.Map;
 import java.util.NoSuchElementException;
 import java.util.Objects;
 import java.util.Optional;
@@ -257,9 +261,10 @@ public void completeHistoryData(String historyDir, String liteDir) {
       //    and return the block numbers of snapshot and history
       BlockNumInfo blockNumInfo = checkAndGetBlockNumInfo(historyDir, liteDir);
       // 2. move archive dbs to bak
-      backupArchiveDbs(liteDir);
+      Map<String, Path> movedDbs = getMovedArchiveDbs(liteDir);
+      backupArchiveDbs(liteDir, movedDbs);
       // 3. copy history data to liteDir
-      copyHistory2Database(historyDir, liteDir);
+      copyHistory2Database(historyDir, liteDir, movedDbs);
       // 4. delete the extra block data in history data
       trimExtraHistory(liteDir, blockNumInfo);
       // 5. merge bak to database
@@ -554,21 +559,86 @@ private BlockNumInfo checkAndGetBlockNumInfo(String historyDir, String liteDir)
     return new BlockNumInfo(snapshotMinNum, historyMaxNum, snapshotMaxNum);
   }
 
-  private void backupArchiveDbs(String databaseDir) throws IOException {
+  @VisibleForTesting
+  static Map<String, Path> getMovedArchiveDbs(String databaseDir) {
+    Map<String, Path> movedDbs = Maps.newHashMap();
+    try {
+      for (String db : archiveDbs) {
+        Path path = Paths.get(databaseDir, db);
+        if (Files.isSymbolicLink(path)) {
+          movedDbs.put(db, path.toRealPath());
+        }
+      }
+    } catch (IOException e) {
+      throw new UncheckedIOException(e);
+    }
+    return movedDbs;
+  }
+
+  private void backupArchiveDbs(String databaseDir, Map<String, Path> movedDbs)
+      throws IOException {
     Path bakDir = Paths.get(databaseDir, BACKUP_DIR_PREFIX + START_TIME);
     logger.info("Backup the archive dbs to {}.", bakDir);
     spec.commandLine().getOut().format("Backup the archive dbs to %s.", bakDir).println();
     if (!FileUtils.createDirIfNotExists(bakDir.toString())) {
       throw new RuntimeException(String.format("create bak dir %s failed", bakDir));
     }
     FileUtils.copyDatabases(Paths.get(databaseDir), bakDir, archiveDbs);
-    archiveDbs.forEach(db -> FileUtils.deleteDir(new File(databaseDir, db)));
+    for (String db : archiveDbs) {
+      Path moved = movedDbs.get(db);
+      if (moved == null) {
+        FileUtils.deleteDir(new File(databaseDir, db));
+      } else {
+        clearDir(moved);
+      }
+    }
   }
 
-  private void copyHistory2Database(String historyDir, String databaseDir) throws IOException {
+  private void copyHistory2Database(String historyDir, String databaseDir,
+      Map<String, Path> movedDbs) throws IOException {
     logger.info("Begin to copy history to database.");
     spec.commandLine().getOut().println("Begin to copy history to database.");
-    FileUtils.copyDatabases(Paths.get(historyDir), Paths.get(databaseDir), archiveDbs);
+    List<String> localDbs = archiveDbs.stream()
+        .filter(db -> !movedDbs.containsKey(db))
+        .collect(Collectors.toList());
+    FileUtils.copyDatabases(Paths.get(historyDir), Paths.get(databaseDir), localDbs);
+    for (Map.Entry<String, Path> moved : movedDbs.entrySet()) {
+      copyInto(Paths.get(historyDir, moved.getKey()), moved.getValue());
+    }
+  }
+
+  @VisibleForTesting
+  static void clearDir(Path dir) {
+    try (Stream<Path> walk = Files.walk(dir)) {
+      List<Path> entries = walk.filter(entry -> !entry.equals(dir))
+          .sorted(Comparator.reverseOrder())
+          .collect(Collectors.toList());
+      for (Path entry : entries) {
+        Files.delete(entry);
+      }
+    } catch (IOException e) {
+      throw new UncheckedIOException(e);
+    }
+  }
+
+  @VisibleForTesting
+  static void copyInto(Path source, Path target) {
+    if (!Files.exists(source)) {
+      return;
+    }
+    try (Stream<Path> walk = Files.walk(source, FileVisitOption.FOLLOW_LINKS)) {
+      List<Path> entries = walk.collect(Collectors.toList());
+      for (Path entry : entries) {
+        Path dest = target.resolve(source.relativize(entry));
+        if (Files.isDirectory(entry)) {
+          Files.createDirectories(dest);
+        } else {
+          FileUtils.copy(entry, dest);
+        }
+      }
+    } catch (IOException e) {
+      throw new UncheckedIOException(e);
+    }
   }
 
   private void trimExtraHistory(String liteDir, BlockNumInfo blockNumInfo)

plugins/src/test/java/org/tron/plugins/DbLiteMergeFailureTest.java

@@ -0,0 +1,145 @@
+package org.tron.plugins;
+
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.StandardCopyOption;
+import java.util.List;
+import org.junit.After;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+import org.mockito.ArgumentMatchers;
+import org.mockito.MockedStatic;
+import org.mockito.Mockito;
+import org.rocksdb.RocksDBException;
+import org.tron.plugins.utils.db.DBInterface;
+import org.tron.plugins.utils.db.DbTool;
+
+public class DbLiteMergeFailureTest {
+
+  @Rule
+  public final TemporaryFolder folder = new TemporaryFolder();
+
+  @After
+  public void closeDatabases() {
+    DbTool.close();
+  }
+
+  @Test
+  public void testBackupFailureLeavesOriginalDbs() throws IOException, RocksDBException {
+    DbLiteMergeFixture fixture = movedFixture();
+    byte[] originalInfo = Files.readAllBytes(fixture.lite.resolve("info.properties"));
+    Path current = fixture.lite.resolve("block").resolve("CURRENT");
+    IOException failure = new IOException("Injected archive backup copy failure");
+
+    try (MockedStatic<Files> files = Mockito.mockStatic(Files.class, Mockito.CALLS_REAL_METHODS)) {
+      files.when(() -> Files.copy(ArgumentMatchers.eq(current), ArgumentMatchers.any(Path.class),
+          ArgumentMatchers.eq(StandardCopyOption.REPLACE_EXISTING))).thenThrow(failure);
+      assertMergeFailed(fixture, originalInfo, failure);
+      files.verify(() -> Files.copy(ArgumentMatchers.eq(current), ArgumentMatchers.any(Path.class),
+          ArgumentMatchers.eq(StandardCopyOption.REPLACE_EXISTING)));
+    }
+
+    assertSnapshotData(fixture.movedRoot);
+    for (String db : DbLiteMergeFixture.archiveDbs()) {
+      assertTrue(db, Files.exists(fixture.movedDirectory(db)
+          .resolve(DbLiteMergeFixture.STALE_FILE)));
+    }
+  }
+
+  @Test
+  public void testClearFailureKeepsRecoverableBackup()
+      throws IOException, RocksDBException {
+    DbLiteMergeFixture fixture = movedFixture();
+    byte[] originalInfo = Files.readAllBytes(fixture.lite.resolve("info.properties"));
+    Path current = fixture.movedDirectory("trans").resolve("CURRENT");
+    IOException failure = new IOException("Injected moved archive deletion failure");
+
+    try (MockedStatic<Files> files = Mockito.mockStatic(Files.class, Mockito.CALLS_REAL_METHODS)) {
+      files.when(() -> Files.delete(ArgumentMatchers.eq(current))).thenThrow(failure);
+      assertMergeFailed(fixture, originalInfo, failure);
+      files.verify(() -> Files.delete(current));
+    }
+
+    assertTrue(Files.exists(current));
+    assertFalse(Files.exists(fixture.movedDirectory("block")
+        .resolve(DbLiteMergeFixture.STALE_FILE)));
+    assertFalse(Files.exists(fixture.movedDirectory("block-index")
+        .resolve(DbLiteMergeFixture.STALE_FILE)));
+    assertRecoverableBackup(fixture);
+  }
+
+  @Test
+  public void testHistoryCopyFailureKeepsRecoverableBackup()
+      throws IOException, RocksDBException {
+    DbLiteMergeFixture fixture = movedFixture();
+    byte[] originalInfo = Files.readAllBytes(fixture.lite.resolve("info.properties"));
+    Path target = fixture.movedDirectory("block");
+    IOException failure = new IOException("Injected history directory creation failure");
+
+    try (MockedStatic<Files> files = Mockito.mockStatic(Files.class, Mockito.CALLS_REAL_METHODS)) {
+      files.when(() -> Files.createDirectories(ArgumentMatchers.eq(target))).thenThrow(failure);
+      assertMergeFailed(fixture, originalInfo, failure);
+      files.verify(() -> Files.createDirectories(target));
+    }
+
+    for (String db : DbLiteMergeFixture.archiveDbs()) {
+      assertFalse(db, Files.exists(fixture.movedDirectory(db)
+          .resolve(DbLiteMergeFixture.STALE_FILE)));
+    }
+    assertRecoverableBackup(fixture);
+  }
+
+  private DbLiteMergeFixture movedFixture() throws IOException, RocksDBException {
+    DbLiteMergeFixture fixture = new DbLiteMergeFixture(folder,
+        DbLiteMergeFixture.SNAPSHOT_MAX - 1);
+    fixture.migrate(DbLiteMergeFixture.archiveDbs());
+    return fixture;
+  }
+
+  private static void assertMergeFailed(DbLiteMergeFixture fixture, byte[] originalInfo,
+      IOException failure) throws IOException {
+    int exitCode = fixture.executeMerge();
+    assertEquals(fixture.error(), 1, exitCode);
+    assertTrue(fixture.error(), fixture.error().contains(failure.getMessage()));
+    fixture.assertLinks();
+    assertArrayEquals(originalInfo, Files.readAllBytes(fixture.lite.resolve("info.properties")));
+    assertEquals(1, fixture.backupDirectories().size());
+  }
+
+  private static void assertRecoverableBackup(DbLiteMergeFixture fixture)
+      throws IOException, RocksDBException {
+    List<Path> backups = fixture.backupDirectories();
+    assertEquals(1, backups.size());
+    assertSnapshotData(backups.get(0));
+    for (String db : DbLiteMergeFixture.archiveDbs()) {
+      assertTrue(db, Files.exists(backups.get(0).resolve(db)
+          .resolve(DbLiteMergeFixture.STALE_FILE)));
+    }
+  }
+
+  private static void assertSnapshotData(Path database) throws IOException, RocksDBException {
+    for (String db : DbLiteMergeFixture.archiveDbs()) {
+      DBInterface store = DbTool.getDB(database.toString(), db);
+      boolean containsGenesis = "block".equals(db) || "block-index".equals(db);
+      if (containsGenesis) {
+        assertArrayEquals(db + " snapshot genesis", DbLiteMergeFixture.value(db, 0),
+            store.get(DbLiteMergeFixture.key(db, 0)));
+      }
+      for (long number = DbLiteMergeFixture.SNAPSHOT_MIN;
+          number <= DbLiteMergeFixture.SNAPSHOT_MAX; number++) {
+        assertArrayEquals(db + " snapshot block " + number,
+            DbLiteMergeFixture.value(db, number), store.get(DbLiteMergeFixture.key(db, number)));
+      }
+      long expectedEntries = DbLiteMergeFixture.SNAPSHOT_MAX - DbLiteMergeFixture.SNAPSHOT_MIN + 1;
+      assertEquals(db + " snapshot entry count", expectedEntries + (containsGenesis ? 1 : 0),
+          store.size());
+    }
+  }
+}