implement backup wal files

Author: zhangchiqing

cmd/util/cmd/find-trie-root/cmd.go

@@ -21,7 +21,8 @@ var (
 	flagRootHash          string
 	flagFrom              int
 	flagTo                int
-	flagOutputDir         string
+	flagBackupDir         string
+	flagTrimAsLatestWAL   bool
 )
 
 // find trie root hash from the wal files.
@@ -50,7 +51,9 @@ func init() {
 	Cmd.Flags().IntVar(&flagFrom, "from", 0, "from segment")
 	Cmd.Flags().IntVar(&flagTo, "to", math.MaxInt32, "to segment")
 
-	Cmd.Flags().StringVar(&flagOutputDir, "output-dir", "", "output directory")
+	Cmd.Flags().StringVar(&flagBackupDir, "backup-dir", "", "directory for backup wal files. must be not exist or empty folder. required when --trim-as-latest-wal flag is set to true.")
+
+	Cmd.Flags().BoolVar(&flagTrimAsLatestWAL, "trim-as-latest-wal", false, "trim the wal file to the last record with the target trie root hash")
 }
 
 func run(*cobra.Command, []string) {
@@ -59,26 +62,46 @@ func run(*cobra.Command, []string) {
 		log.Fatal().Err(err).Msg("cannot parse input")
 	}
 
-	if flagExecutionStateDir == flagOutputDir {
-		log.Fatal().Msg("output directory cannot be the same as the execution state directory")
+	if flagExecutionStateDir == flagBackupDir {
+		log.Fatal().Msg("backup directory cannot be the same as the execution state directory")
 	}
 
 	segment, offset, err := searchRootHashInSegments(rootHash, flagExecutionStateDir, flagFrom, flagTo)
 	if err != nil {
 		log.Fatal().Err(err).Msg("cannot find root hash in segments")
 	}
-	log.Info().Msgf("found root hash in segment %d at offset %d", segment, offset)
 
-	if len(flagOutputDir) == 0 {
+	segmentFile := prometheusWAL.SegmentName(flagExecutionStateDir, segment)
+
+	log.Info().Msgf("found root hash in segment %d at offset %d, segment file: %v", segment, offset, segmentFile)
+
+	if !flagTrimAsLatestWAL {
+		log.Info().Msg("not trimming WAL. Exiting. to trim the WAL, use --trim-as-latest-wal flag")
 		return
 	}
 
-	err = copyWAL(flagExecutionStateDir, flagOutputDir, segment, rootHash)
+	if len(flagBackupDir) == 0 {
+		log.Error().Msgf("backup directory is not provided")
+		return
+	}
+
+	// genereate a segment file to the temporary folder with the root hash as its last record
+	newSegmentFile, err := findRootHashAndCreateTrimmed(flagExecutionStateDir, segment, rootHash)
 	if err != nil {
 		log.Fatal().Err(err).Msg("cannot copy WAL")
 	}
 
-	log.Info().Msgf("copied WAL to %s", flagOutputDir)
+	log.Info().Msgf("successfully copied WAL to the temporary folder %v", newSegmentFile)
+
+	// before replacing the last wal file with the newly generated one, backup the rollbacked wals
+	// then move the last segment file to the execution state directory
+	err = backupRollbackedWALsAndMoveLastSegmentFile(segment, flagExecutionStateDir, flagBackupDir, newSegmentFile)
+	if err != nil {
+		log.Fatal().Err(err).Msg("cannot backup rollbacked WALs")
+	}
+
+	log.Info().Msgf("successfully trimmed WAL %v the trie root hash %v as its last record, original wal files are moved to %v",
+		segment, rootHash, flagBackupDir)
 }
 
 func parseInput(rootHashStr string) (ledger.RootHash, error) {
@@ -181,26 +204,33 @@ func searchRootHashInSegments(
 	return 0, 0, fmt.Errorf("finish reading all segment files from %d to %d, but not found", from, to)
 }
 
-func copyWAL(dir, outputDir string, segment int, expectedRoot ledger.RootHash) error {
-	writer, err := prometheusWAL.NewSize(log.Logger, nil, outputDir, wal.SegmentSize, false)
+// findRootHashAndCreateTrimmed finds the root hash in the segment file from the given dir folder
+// and creates a new segment file with the expected root hash as the last record in a temporary folder.
+// it return the path to the new segment file.
+func findRootHashAndCreateTrimmed(dir string, segment int, expectedRoot ledger.RootHash) (string, error) {
+	tmpFolder, err := os.MkdirTemp("", "flow-last-segment-file")
 	if err != nil {
-		return fmt.Errorf("cannot create writer WAL: %w", err)
+		return "", fmt.Errorf("cannot create temporary folder: %w", err)
 	}
 
-	defer writer.Close()
+	newSegmentFile := prometheusWAL.SegmentName(tmpFolder, segment)
+
+	log.Info().Msgf("writing new segment file to %v", newSegmentFile)
 
-	w, err := prometheusWAL.NewSize(log.Logger, nil, dir, wal.SegmentSize, false)
+	writer, err := prometheusWAL.NewSize(log.Logger, nil, tmpFolder, wal.SegmentSize, false)
 	if err != nil {
-		return fmt.Errorf("cannot create WAL: %w", err)
+		return "", fmt.Errorf("cannot create writer WAL: %w", err)
 	}
 
+	defer writer.Close()
+
 	sr, err := prometheusWAL.NewSegmentsRangeReader(log.Logger, prometheusWAL.SegmentRange{
-		Dir:   w.Dir(),
+		Dir:   dir,
 		First: segment,
 		Last:  segment,
 	})
 	if err != nil {
-		return fmt.Errorf("cannot create WAL segments reader: %w", err)
+		return "", fmt.Errorf("cannot create WAL segments reader: %w", err)
 	}
 
 	defer sr.Close()
@@ -211,7 +241,7 @@ func copyWAL(dir, outputDir string, segment int, expectedRoot ledger.RootHash) e
 		record := reader.Record()
 		operation, _, update, err := wal.Decode(record)
 		if err != nil {
-			return fmt.Errorf("cannot decode LedgerWAL record: %w", err)
+			return "", fmt.Errorf("cannot decode LedgerWAL record: %w", err)
 		}
 
 		switch operation {
@@ -220,23 +250,101 @@ func copyWAL(dir, outputDir string, segment int, expectedRoot ledger.RootHash) e
 			bytes := wal.EncodeUpdate(update)
 			_, err = writer.Log(bytes)
 			if err != nil {
-				return fmt.Errorf("cannot write LedgerWAL record: %w", err)
+				return "", fmt.Errorf("cannot write LedgerWAL record: %w", err)
 			}
 
 			rootHash := update.RootHash
 
 			if rootHash.Equals(expectedRoot) {
 				log.Info().Msgf("found expected trie root hash %v, finish writing", rootHash)
-				return nil
+				return newSegmentFile, nil
 			}
 		default:
 		}
 
 		err = reader.Err()
 		if err != nil {
-			return fmt.Errorf("cannot read LedgerWAL: %w", err)
+			return "", fmt.Errorf("cannot read LedgerWAL: %w", err)
 		}
 	}
 
-	return fmt.Errorf("finish reading all segment files from %d to %d, but not found", segment, segment)
+	return "", fmt.Errorf("finish reading all segment files from %d to %d, but not found", segment, segment)
+}
+
+func checkFolderNotExistOrEmpty(folderPath string) (bool, error) {
+	// Check if the folder exists
+	info, err := os.Stat(folderPath)
+	if err != nil {
+		if os.IsNotExist(err) {
+			return true, nil
+		}
+		return false, nil
+	}
+
+	// Check if the path is a directory
+	if !info.IsDir() {
+		return false, fmt.Errorf("The path is not a directory.")
+	}
+
+	// Check if the folder is empty
+	files, err := os.ReadDir(folderPath)
+	if err != nil {
+		return false, fmt.Errorf("Cannot read the folder.")
+	}
+
+	return len(files) == 0, nil
+}
+
+// backup new wals before replacing
+func backupRollbackedWALsAndMoveLastSegmentFile(
+	segment int, walDir, backupDir string, newSegmentFile string) error {
+	// making sure the backup dir is empty
+	empty, err := checkFolderNotExistOrEmpty(backupDir)
+	if err != nil {
+		return fmt.Errorf("cannot check backup directory: %w", err)
+	}
+
+	if !empty {
+		return fmt.Errorf("backup directory %s is not empty", backupDir)
+	}
+
+	// Create the backup directory
+	err = os.MkdirAll(backupDir, os.ModePerm)
+	if err != nil {
+		return fmt.Errorf("cannot create backup directory: %w", err)
+	}
+
+	first, last, err := prometheusWAL.Segments(walDir)
+	if err != nil {
+		return fmt.Errorf("cannot get segments: %w", err)
+	}
+
+	if segment < first {
+		return fmt.Errorf("segment %d is less than the first segment %d", segment, first)
+	}
+
+	// backup all the segment files that have higher number than the given segment, including
+	// the segment file itself, since it will be replaced.
+	for i := segment; i <= last; i++ {
+		segmentFile := prometheusWAL.SegmentName(walDir, i)
+		backupFile := prometheusWAL.SegmentName(backupDir, i)
+
+		log.Info().Msgf("backup segment file %s to %s, %v/%v", segmentFile, backupFile, i, last)
+		err := os.Rename(segmentFile, backupFile)
+		if err != nil {
+			return fmt.Errorf("cannot move segment file %s to %s: %w", segmentFile, backupFile, err)
+		}
+	}
+
+	// after backup the segment files, replace the last segment file
+	segmentToBeReplaced := prometheusWAL.SegmentName(walDir, segment)
+
+	log.Info().Msgf("moving segment file %s to %s", newSegmentFile, segmentToBeReplaced)
+
+	err = os.Rename(newSegmentFile, segmentToBeReplaced)
+	if err != nil {
+		return fmt.Errorf("cannot move segment file %s to %s: %w", newSegmentFile, segmentToBeReplaced, err)
+	}
+
+	return nil
 }