Skip to content

Commit

Permalink
fix(gui): resolve reset and commit issues for disk code cache
Browse files Browse the repository at this point in the history
  • Loading branch information
skylot committed Dec 16, 2023
1 parent b6155af commit 8db70ee
Show file tree
Hide file tree
Showing 7 changed files with 153 additions and 114 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@

public class ReloadProject implements IJadxEvent {

public static final ReloadProject INSTANCE = new ReloadProject();
public static final ReloadProject EVENT = new ReloadProject();

private ReloadProject() {
// singleton
Expand Down
221 changes: 120 additions & 101 deletions jadx-gui/src/main/java/jadx/gui/cache/code/disk/DiskCodeCache.java
Original file line number Diff line number Diff line change
@@ -1,26 +1,21 @@
package jadx.gui.cache.code.disk;

import java.io.BufferedInputStream;
import java.io.BufferedOutputStream;
import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.BitSet;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.stream.Stream;

import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
Expand All @@ -32,44 +27,37 @@
import jadx.core.Jadx;
import jadx.core.dex.nodes.ClassNode;
import jadx.core.dex.nodes.RootNode;
import jadx.core.utils.StringUtils;
import jadx.core.utils.Utils;
import jadx.core.utils.exceptions.JadxRuntimeException;
import jadx.core.utils.files.FileUtils;

import static java.nio.file.StandardOpenOption.CREATE;
import static java.nio.file.StandardOpenOption.TRUNCATE_EXISTING;
import static java.nio.file.StandardOpenOption.WRITE;

public class DiskCodeCache implements ICodeCache {
private static final Logger LOG = LoggerFactory.getLogger(DiskCodeCache.class);

private static final int DATA_FORMAT_VERSION = 13;

private static final byte[] JADX_NAMES_MAP_HEADER = "jadxnm".getBytes(StandardCharsets.US_ASCII);
private static final int DATA_FORMAT_VERSION = 14;

private final Path baseDir;
private final Path srcDir;
private final Path metaDir;
private final Path codeVersionFile;
private final Path namesMapFile;
private final String codeVersion;
private final CodeMetadataAdapter codeMetadataAdapter;
private final ExecutorService writePool;
private final Map<String, ICodeInfo> writeOps = new ConcurrentHashMap<>();
private final Map<String, Integer> namesMap = new ConcurrentHashMap<>();
private final Map<String, Integer> allClsIds;
private final Map<String, CacheData> clsDataMap;

public DiskCodeCache(RootNode root, Path baseDir) {
public DiskCodeCache(RootNode root, Path projectCacheDir) {
baseDir = projectCacheDir.resolve("code");
srcDir = baseDir.resolve("sources");
metaDir = baseDir.resolve("metadata");
codeVersionFile = baseDir.resolve("code-version");
namesMapFile = baseDir.resolve("names-map");
JadxArgs args = root.getArgs();
codeVersion = buildCodeVersion(args, root.getDecompiler());
writePool = Executors.newFixedThreadPool(args.getThreadsCount());
codeMetadataAdapter = new CodeMetadataAdapter(root);
allClsIds = buildClassIdsMap(root.getClasses());
clsDataMap = buildClassDataMap(root.getClasses());
if (checkCodeVersion()) {
loadNamesMap();
loadCachedSet();
} else {
reset();
}
Expand All @@ -91,10 +79,12 @@ private boolean checkCodeVersion() {
private void reset() {
try {
long start = System.currentTimeMillis();
LOG.info("Resetting disk code cache, base dir: {}", srcDir.getParent().toAbsolutePath());
FileUtils.deleteDirIfExists(srcDir);
FileUtils.deleteDirIfExists(metaDir);
FileUtils.deleteFileIfExists(namesMapFile);
LOG.info("Resetting disk code cache, base dir: {}", baseDir.toAbsolutePath());
FileUtils.deleteDirIfExists(baseDir);
if (Files.exists(baseDir.getParent().resolve(codeVersionFile.getFileName()))) {
// remove old version cache files
FileUtils.deleteDirIfExists(baseDir.getParent());
}
FileUtils.makeDirs(srcDir);
FileUtils.makeDirs(metaDir);
FileUtils.writeFile(codeVersionFile, codeVersion);
Expand All @@ -104,7 +94,7 @@ private void reset() {
} catch (Exception e) {
throw new JadxRuntimeException("Failed to reset code cache", e);
} finally {
namesMap.clear();
clsDataMap.values().forEach(d -> d.setCached(false));
}
}

Expand All @@ -113,18 +103,22 @@ private void reset() {
*/
@Override
public void add(String clsFullName, ICodeInfo codeInfo) {
writeOps.put(clsFullName, codeInfo);
int clsId = getClsId(clsFullName);
namesMap.put(clsFullName, clsId);
CacheData clsData = getClsData(clsFullName);
clsData.setTmpCodeInfo(codeInfo);
clsData.setCached(true);
writePool.execute(() -> {
try {
FileUtils.writeFile(getJavaFile(clsId), codeInfo.getCodeStr());
codeMetadataAdapter.write(getMetadataFile(clsId), codeInfo.getCodeMetadata());
int clsId = clsData.getClsId();
ICodeInfo code = clsData.getTmpCodeInfo();
if (code != null) {
FileUtils.writeFile(getJavaFile(clsId), code.getCodeStr());
codeMetadataAdapter.write(getMetadataFile(clsId), code.getCodeMetadata());
}
} catch (Exception e) {
LOG.error("Failed to write code cache for " + clsFullName, e);
remove(clsFullName);
} finally {
writeOps.remove(clsFullName);
clsData.setTmpCodeInfo(null);
}
});
}
Expand All @@ -135,12 +129,12 @@ public void add(String clsFullName, ICodeInfo codeInfo) {
if (!contains(clsFullName)) {
return null;
}
ICodeInfo wrtCodeInfo = writeOps.get(clsFullName);
if (wrtCodeInfo != null) {
return wrtCodeInfo.getCodeStr();
CacheData clsData = getClsData(clsFullName);
ICodeInfo tmpCodeInfo = clsData.getTmpCodeInfo();
if (tmpCodeInfo != null) {
return tmpCodeInfo.getCodeStr();
}
int clsId = getClsId(clsFullName);
Path javaFile = getJavaFile(clsId);
Path javaFile = getJavaFile(clsData.getClsId());
if (!Files.exists(javaFile)) {
return null;
}
Expand All @@ -152,16 +146,17 @@ public void add(String clsFullName, ICodeInfo codeInfo) {
}

@Override
public ICodeInfo get(String clsFullName) {
public @NotNull ICodeInfo get(String clsFullName) {
try {
if (!contains(clsFullName)) {
return ICodeInfo.EMPTY;
}
ICodeInfo wrtCodeInfo = writeOps.get(clsFullName);
if (wrtCodeInfo != null) {
return wrtCodeInfo;
CacheData clsData = getClsData(clsFullName);
ICodeInfo tmpCodeInfo = clsData.getTmpCodeInfo();
if (tmpCodeInfo != null) {
return tmpCodeInfo;
}
int clsId = getClsId(clsFullName);
int clsId = clsData.getClsId();
Path javaFile = getJavaFile(clsId);
if (!Files.exists(javaFile)) {
return ICodeInfo.EMPTY;
Expand All @@ -176,17 +171,24 @@ public ICodeInfo get(String clsFullName) {

@Override
public boolean contains(String clsFullName) {
return namesMap.containsKey(clsFullName);
return getClsData(clsFullName).isCached();
}

@Override
public void remove(String clsFullName) {
try {
Integer clsId = namesMap.remove(clsFullName);
if (clsId != null) {
LOG.debug("Removing class info from disk: {}", clsFullName);
Files.deleteIfExists(getJavaFile(clsId));
Files.deleteIfExists(getMetadataFile(clsId));
CacheData clsData = getClsData(clsFullName);
if (clsData.isCached()) {
clsData.setCached(false);
if (clsData.getTmpCodeInfo() == null) {
LOG.debug("Removing class info from disk: {}", clsFullName);
int clsId = clsData.getClsId();
Files.deleteIfExists(getJavaFile(clsId));
Files.deleteIfExists(getMetadataFile(clsId));
} else {
// class info not yet written to disk
clsData.setTmpCodeInfo(null);
}
}
} catch (Exception e) {
throw new JadxRuntimeException("Failed to remove code cache for " + clsFullName, e);
Expand All @@ -206,55 +208,39 @@ private String buildCodeVersion(JadxArgs args, @Nullable JadxDecompiler decompil
+ ":" + FileUtils.buildInputsHash(Utils.collectionMap(inputFiles, File::toPath));
}

private int getClsId(String clsFullName) {
Integer id = allClsIds.get(clsFullName);
if (id == null) {
private CacheData getClsData(String clsFullName) {
CacheData clsData = clsDataMap.get(clsFullName);
if (clsData == null) {
throw new JadxRuntimeException("Unknown class name: " + clsFullName);
}
return id;
return clsData;
}

private void saveNamesMap() {
LOG.debug("Saving names map for disk cache...");
try (OutputStream fileOutput = Files.newOutputStream(namesMapFile, WRITE, CREATE, TRUNCATE_EXISTING);
DataOutputStream out = new DataOutputStream(new BufferedOutputStream(fileOutput))) {
out.write(JADX_NAMES_MAP_HEADER);
out.writeInt(namesMap.size());
for (Map.Entry<String, Integer> entry : namesMap.entrySet()) {
out.writeUTF(entry.getKey());
out.writeInt(entry.getValue());
}
private void loadCachedSet() {
long start = System.currentTimeMillis();
BitSet cachedSet = new BitSet(clsDataMap.size());
try (Stream<Path> stream = Files.walk(metaDir)) {
stream.forEach(file -> {
String fileName = file.getFileName().toString();
if (fileName.endsWith(".jadxmd")) {
String idStr = StringUtils.removeSuffix(fileName, ".jadxmd");
int clsId = Integer.parseInt(idStr, 16);
cachedSet.set(clsId);
}
});
} catch (Exception e) {
throw new JadxRuntimeException("Failed to save names map file", e);
throw new JadxRuntimeException("Failed to enumerate cached classes", e);
}
}

private void loadNamesMap() {
if (!Files.exists(namesMapFile)) {
reset();
return;
}
namesMap.clear();
try (InputStream fileInput = Files.newInputStream(namesMapFile);
DataInputStream in = new DataInputStream(new BufferedInputStream(fileInput))) {
in.skipBytes(JADX_NAMES_MAP_HEADER.length);
int count = in.readInt();
for (int i = 0; i < count; i++) {
String clsName = in.readUTF();
int clsId = in.readInt();
namesMap.put(clsName, clsId);
Integer prevId = allClsIds.get(clsName);
if (prevId == null || prevId != clsId) {
LOG.debug("Unexpected class id, got: {}, expect: {}", clsId, prevId);
LOG.warn("Inconsistent disk cache, resetting...");
reset();
return;
}
int count = 0;
for (CacheData data : clsDataMap.values()) {
int clsId = data.getClsId();
if (cachedSet.get(clsId)) {
data.setCached(true);
count++;
}
LOG.info("Found {} classes in disk cache, dir: {}", count, metaDir.getParent());
} catch (Exception e) {
throw new JadxRuntimeException("Failed to load names map file", e);
}
LOG.info("Found {} classes in disk cache, time: {}ms, dir: {}",
count, System.currentTimeMillis() - start, metaDir.getParent());
}

private Path getJavaFile(int clsId) {
Expand All @@ -271,25 +257,58 @@ private Path getPathForClsId(int clsId, String ext) {
return Paths.get(firstByte, FileUtils.intToHex(clsId) + ext);
}

private Map<String, Integer> buildClassIdsMap(List<ClassNode> classes) {
private Map<String, CacheData> buildClassDataMap(List<ClassNode> classes) {
int clsCount = classes.size();
Map<String, Integer> map = new HashMap<>(clsCount);
Map<String, CacheData> map = new HashMap<>(clsCount);
for (int i = 0; i < clsCount; i++) {
ClassNode cls = classes.get(i);
map.put(cls.getRawName(), i);
map.put(cls.getRawName(), new CacheData(i));
}
return map;
}

@SuppressWarnings("ResultOfMethodCallIgnored")
@Override
public void close() throws IOException {
try {
saveNamesMap();
writePool.shutdown();
writePool.awaitTermination(2, TimeUnit.MINUTES);
} catch (InterruptedException e) {
LOG.error("Failed to finish file writes", e);
synchronized (this) {
try {
writePool.shutdown();
boolean completed = writePool.awaitTermination(1, TimeUnit.MINUTES);
if (!completed) {
LOG.warn("Disk code cache closing terminated by timeout");
}
} catch (InterruptedException e) {
LOG.error("Failed to close disk code cache", e);
}
}
}

private static final class CacheData {
private final int clsId;
private boolean cached;
private @Nullable ICodeInfo tmpCodeInfo;

public CacheData(int clsId) {
this.clsId = clsId;
}

public int getClsId() {
return clsId;
}

public boolean isCached() {
return cached;
}

public void setCached(boolean cached) {
this.cached = cached;
}

public @Nullable ICodeInfo getTmpCodeInfo() {
return tmpCodeInfo;
}

public void setTmpCodeInfo(@Nullable ICodeInfo tmpCodeInfo) {
this.tmpCodeInfo = tmpCodeInfo;
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -111,7 +111,7 @@ private Path buildCacheDir(JadxProject project) {
if (Objects.equals(cacheDirValue, ".")) {
return buildLocalCacheDir(project);
}
Path cacheBaseDir = cacheDirValue == null ? JadxFiles.CACHE_DIR : Paths.get(cacheDirValue);
Path cacheBaseDir = cacheDirValue == null ? JadxFiles.PROJECTS_CACHE_DIR : Paths.get(cacheDirValue);
return cacheBaseDir.resolve(buildProjectUniqName(project));
}

Expand Down
Loading

0 comments on commit 8db70ee

Please sign in to comment.