Skip to content

Commit 982425f

Browse files
halibobo1205halibobo1205
committed
feat(toolkit): add RocksDB compact command
1 parent 4be5481 commit 982425f

4 files changed

Lines changed: 230 additions & 0 deletions

File tree

‎plugins/src/main/java/common/org/tron/plugins/Db.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
DbMove.class,
1111
DbArchive.class,
1212
DbConvert.class,
13+
DbCompact.class,
1314
DbLite.class,
1415
DbCopy.class,
1516
DbRoot.class

‎plugins/src/main/java/common/org/tron/plugins/DbCompact.java‎

Lines changed: 178 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,178 @@
1+
package org.tron.plugins;
2+
3+
import java.io.File;
4+
import java.nio.file.Path;
5+
import java.nio.file.Paths;
6+
import java.util.List;
7+
import java.util.Objects;
8+
import java.util.concurrent.Callable;
9+
import java.util.stream.Collectors;
10+
import lombok.extern.slf4j.Slf4j;
11+
import me.tongfei.progressbar.ProgressBar;
12+
import org.rocksdb.Options;
13+
import org.rocksdb.RocksDB;
14+
import org.tron.plugins.utils.DBUtils;
15+
import org.tron.plugins.utils.FileUtils;
16+
import picocli.CommandLine;
17+
18+
@Slf4j(topic = "compact")
19+
@CommandLine.Command(name = "compact",
20+
description = "Compact one or more RocksDB databases.",
21+
exitCodeListHeading = "Exit Codes:%n",
22+
exitCodeList = {
23+
"0:Successful",
24+
"n:Internal error: exception occurred,please check toolkit.log"})
25+
public class DbCompact implements Callable<Integer> {
26+
27+
static {
28+
RocksDB.loadLibrary();
29+
}
30+
31+
@CommandLine.Spec
32+
CommandLine.Model.CommandSpec spec;
33+
34+
@CommandLine.Parameters(index = "0", defaultValue = "output-directory/database",
35+
description = "Input database parent path. Default: ${DEFAULT-VALUE}")
36+
private Path db;
37+
38+
@CommandLine.Option(names = {"--dbs"}, split = ",",
39+
description = "Database names to compact. Supports comma-separated values, "
40+
+ "for example: --dbs account-asset,account")
41+
private List<String> dbs;
42+
43+
@CommandLine.Option(names = {"-h", "--help"}, help = true, description = "display a help message")
44+
private boolean help;
45+
46+
@Override
47+
public Integer call() {
48+
if (help) {
49+
spec.commandLine().usage(System.out);
50+
return 0;
51+
}
52+
if (!db.toFile().exists()) {
53+
logger.info("{} does not exist.", db);
54+
spec.commandLine().getErr().println(spec.commandLine().getColorScheme()
55+
.errorText(String.format("%s does not exist.", db)));
56+
return 404;
57+
}
58+
if (dbs == null || dbs.isEmpty()) {
59+
String tips = "Specify at least one database: --dbs dbName[,dbName...].";
60+
logger.info(tips);
61+
spec.commandLine().getErr().println(spec.commandLine().getColorScheme().errorText(tips));
62+
return 404;
63+
}
64+
65+
List<CompactTask> tasks = dbs.stream()
66+
.map(String::trim)
67+
.filter(name -> !name.isEmpty())
68+
.distinct()
69+
.map(name -> new CompactTask(db, name))
70+
.collect(Collectors.toList());
71+
if (tasks.isEmpty()) {
72+
String tips = "Specify at least one non-empty database name.";
73+
logger.info(tips);
74+
spec.commandLine().getErr().println(spec.commandLine().getColorScheme().errorText(tips));
75+
return 404;
76+
}
77+
78+
long start = System.currentTimeMillis();
79+
List<CompactResult> results = ProgressBar.wrap(tasks.stream(), "compact task")
80+
.parallel()
81+
.map(CompactTask::compact)
82+
.collect(Collectors.toList());
83+
results.forEach(this::printInfo);
84+
85+
List<String> fails = results.stream()
86+
.filter(result -> !result.success)
87+
.map(result -> result.name)
88+
.collect(Collectors.toList());
89+
long during = (System.currentTimeMillis() - start) / 1000;
90+
spec.commandLine().getOut().format("compact db done, fails: %s, take %d s.",
91+
fails, during).println();
92+
logger.info("database compact use {} seconds total, fails: {}.", during, fails);
93+
return fails.size();
94+
}
95+
96+
private void printInfo(CompactResult result) {
97+
if (result.success) {
98+
spec.commandLine().getOut().println(result.message);
99+
} else {
100+
spec.commandLine().getErr().println(spec.commandLine().getColorScheme()
101+
.errorText(result.message));
102+
}
103+
}
104+
105+
private static class CompactTask {
106+
private final Path parent;
107+
private final String name;
108+
private final Path path;
109+
110+
private CompactTask(Path parent, String name) {
111+
this.parent = parent;
112+
this.name = name;
113+
this.path = Paths.get(parent.toString(), name);
114+
}
115+
116+
private CompactResult compact() {
117+
long start = System.currentTimeMillis();
118+
try {
119+
if (!path.toFile().exists()) {
120+
return CompactResult.fail(name, String.format("db: %s, fail: %s does not exist.",
121+
name, path));
122+
}
123+
if (!isRocksDb(path)) {
124+
return CompactResult.fail(name, String.format("db: %s, fail: not a RocksDB database.",
125+
name));
126+
}
127+
if (DBUtils.MARKET_PAIR_PRICE_TO_ORDER.equalsIgnoreCase(name)) {
128+
return CompactResult.ok(name, String.format("db: %s, skipped.", name));
129+
}
130+
131+
logger.info("compact database {} start", name);
132+
try (Options options = DBUtils.newDefaultRocksDbOptions(false, name)) {
133+
options.setCreateIfMissing(false);
134+
try (RocksDB rocks = RocksDB.open(options, path.toString())) {
135+
rocks.compactRange();
136+
}
137+
}
138+
long during = (System.currentTimeMillis() - start) / 1000;
139+
logger.info("compact database {} end, take {} s", name, during);
140+
return CompactResult.ok(name, String.format("db: %s, compacted, take %d s.",
141+
name, during));
142+
} catch (Exception e) {
143+
logger.error("compact database {} fail", name, e);
144+
return CompactResult.fail(name, String.format("db: %s, fail: %s",
145+
name, e.getMessage()));
146+
}
147+
}
148+
149+
private static boolean isRocksDb(Path path) {
150+
String engineFile = path + File.separator + DBUtils.FILE_ENGINE;
151+
if (!FileUtils.isExists(engineFile)) {
152+
return false;
153+
}
154+
String engine = FileUtils.readProperty(engineFile, DBUtils.KEY_ENGINE);
155+
return DBUtils.ROCKSDB.equalsIgnoreCase(engine);
156+
}
157+
}
158+
159+
private static class CompactResult {
160+
private final String name;
161+
private final boolean success;
162+
private final String message;
163+
164+
private CompactResult(String name, boolean success, String message) {
165+
this.name = name;
166+
this.success = success;
167+
this.message = Objects.toString(message, "");
168+
}
169+
170+
private static CompactResult ok(String name, String message) {
171+
return new CompactResult(name, true, message);
172+
}
173+
174+
private static CompactResult fail(String name, String message) {
175+
return new CompactResult(name, false, message);
176+
}
177+
}
178+
}

‎plugins/src/main/java/common/org/tron/plugins/utils/DBUtils.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
import org.tron.common.arch.Arch;
1818
import org.tron.common.utils.MarketOrderPriceComparatorForLevelDB;
1919
import org.tron.common.utils.MarketOrderPriceComparatorForRocksDB;
20+
import org.tron.common.utils.RocksDbPartitionedFilter;
2021
import org.tron.protos.Protocol;
2122

2223
public class DBUtils {
@@ -123,6 +124,7 @@ public static Options newDefaultRocksDbOptions(boolean forBulkLoad, String name)
123124
tableCfg.setCacheIndexAndFilterBlocks(true);
124125
tableCfg.setPinL0FilterAndIndexBlocksInCache(true);
125126
tableCfg.setFilter(new BloomFilter(10, false));
127+
RocksDbPartitionedFilter.enable(tableCfg);
126128
options.setTableFormatConfig(tableCfg);
127129
if (forBulkLoad) {
128130
options.prepareForBulkLoad();

‎plugins/src/test/java/org/tron/plugins/rocksdb/DbCompactTest.java‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
package org.tron.plugins.rocksdb;
2+
3+
import java.io.IOException;
4+
import java.util.UUID;
5+
import org.junit.Assert;
6+
import org.junit.Test;
7+
import org.rocksdb.RocksDBException;
8+
import org.tron.plugins.DbTest;
9+
import org.tron.plugins.Toolkit;
10+
import org.tron.plugins.utils.db.DbTool;
11+
import picocli.CommandLine;
12+
13+
public class DbCompactTest extends DbTest {
14+
15+
@Test
16+
public void testRunForRocksDB() throws IOException, RocksDBException {
17+
init(DbTool.DbType.RocksDB);
18+
String[] args = new String[] {"db", "compact", INPUT_DIRECTORY,
19+
"--dbs", "account,market_pair_price_to_order"};
20+
Assert.assertEquals(0, cli.execute(args));
21+
}
22+
23+
@Test
24+
public void testHelp() {
25+
String[] args = new String[] {"db", "compact", "-h"};
26+
CommandLine cli = new CommandLine(new Toolkit());
27+
Assert.assertEquals(0, cli.execute(args));
28+
}
29+
30+
@Test
31+
public void testNotExist() {
32+
String[] args = new String[] {"db", "compact", UUID.randomUUID().toString(),
33+
"--dbs", "account"};
34+
Assert.assertEquals(404, cli.execute(args));
35+
}
36+
37+
@Test
38+
public void testEmptyDbs() throws IOException {
39+
String[] args = new String[] {"db", "compact", temporaryFolder.newFolder().toString()};
40+
Assert.assertEquals(404, cli.execute(args));
41+
}
42+
43+
@Test
44+
public void testMissingDb() throws IOException {
45+
String[] args = new String[] {"db", "compact", temporaryFolder.newFolder().toString(),
46+
"--dbs", "missing"};
47+
Assert.assertEquals(1, cli.execute(args));
48+
}
49+
}

0 commit comments

Comments
 (0)