jessikitty/jadx
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

refactor: add new task executor, remove task barrier (#1879)

Browse Source
This commit is contained in:
Skylot
2024-01-23 20:45:19 +00:00
parent e73612b4d2
commit 75d2e540aa
15 changed files with 398 additions and 192 deletions
Download Patch File Download Diff File Expand all files Collapse all files
jadx-core/src/main/java/jadx/api/JadxDecompiler.java
+73 -83
View File
@@ -12,9 +12,6 @@ import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.function.Predicate;
import java.util.stream.Collectors;
@@ -36,6 +33,7 @@ import jadx.api.plugins.input.JadxCodeInput;
import jadx.api.plugins.pass.JadxPass;
import jadx.api.plugins.pass.types.JadxAfterLoadPass;
import jadx.api.plugins.pass.types.JadxPassType;
import jadx.api.utils.tasks.ITaskExecutor;
import jadx.core.Jadx;
import jadx.core.dex.attributes.AFlag;
import jadx.core.dex.nodes.ClassNode;
@@ -52,6 +50,7 @@ import jadx.core.utils.DecompilerScheduler;
import jadx.core.utils.Utils;
import jadx.core.utils.exceptions.JadxRuntimeException;
import jadx.core.utils.files.FileUtils;
import jadx.core.utils.tasks.TaskExecutor;
import jadx.core.xmlgen.BinaryXMLParser;
import jadx.core.xmlgen.ProtoXMLParser;
import jadx.core.xmlgen.ResourcesSaver;
@@ -248,13 +247,12 @@ public final class JadxDecompiler implements Closeable {
@SuppressWarnings("BusyWait")
public void save(int intervalInMillis, ProgressListener listener) {
ThreadPoolExecutor ex = (ThreadPoolExecutor) getSaveExecutor();
ex.shutdown();
try {
long total = ex.getTaskCount();
while (ex.isTerminating()) {
long done = ex.getCompletedTaskCount();
listener.progress(done, total);
ITaskExecutor tasks = getSaveTaskExecutor();
tasks.execute();
long total = tasks.getTasksCount();
while (tasks.isRunning()) {
listener.progress(tasks.getProgress(), total);
Thread.sleep(intervalInMillis);
}
} catch (InterruptedException e) {
@@ -271,72 +269,64 @@ public final class JadxDecompiler implements Closeable {
save(false, true);
}
@SuppressWarnings("ResultOfMethodCallIgnored")
private void save(boolean saveSources, boolean saveResources) {
ExecutorService ex = getSaveExecutor(saveSources, saveResources);
ex.shutdown();
try {
ex.awaitTermination(1, TimeUnit.DAYS);
} catch (InterruptedException e) {
LOG.error("Save interrupted", e);
Thread.currentThread().interrupt();
}
ITaskExecutor executor = getSaveTasks(saveSources, saveResources);
executor.execute();
executor.awaitTermination();
}
public ExecutorService getSaveExecutor() {
return getSaveExecutor(!args.isSkipSources(), !args.isSkipResources());
}
public List<Runnable> getSaveTasks() {
public ITaskExecutor getSaveTaskExecutor() {
return getSaveTasks(!args.isSkipSources(), !args.isSkipResources());
}
private ExecutorService getSaveExecutor(boolean saveSources, boolean saveResources) {
int threadsCount = args.getThreadsCount();
LOG.debug("processing threads count: {}", threadsCount);
LOG.info("processing ...");
ExecutorService executor = Executors.newFixedThreadPool(threadsCount);
List<Runnable> tasks = getSaveTasks(saveSources, saveResources);
tasks.forEach(executor::execute);
return executor;
@Deprecated(forRemoval = true)
public ExecutorService getSaveExecutor() {
ITaskExecutor executor = getSaveTaskExecutor();
executor.execute();
return executor.getInternalExecutor();
}
private List<Runnable> getSaveTasks(boolean saveSources, boolean saveResources) {
@Deprecated(forRemoval = true)
public List<Runnable> getSaveTasks() {
return Collections.singletonList(this::save);
}
private TaskExecutor getSaveTasks(boolean saveSources, boolean saveResources) {
if (root == null) {
throw new JadxRuntimeException("No loaded files");
}
File sourcesOutDir;
File resOutDir;
List<Runnable> tasks = new ArrayList<>();
TaskBarrier barrier = new TaskBarrier();
ExportGradleTask gradleExportTask;
if (args.isExportAsGradleProject()) {
ExportGradleTask gradleExportTask = new ExportGradleTask(resources, root, args.getOutDir(), barrier);
gradleExportTask = new ExportGradleTask(resources, root, args.getOutDir());
gradleExportTask.init();
sourcesOutDir = gradleExportTask.getSrcOutDir();
resOutDir = gradleExportTask.getResOutDir();
tasks.add(gradleExportTask);
} else {
sourcesOutDir = args.getOutDirSrc();
resOutDir = args.getOutDirRes();
gradleExportTask = null;
}
int taskCount = 0;
// save resources first because decompilation can hang or fail
TaskExecutor executor = new TaskExecutor();
executor.setThreadsCount(args.getThreadsCount());
if (saveResources) {
taskCount = appendResourcesSaveTasks(tasks, resOutDir, barrier);
// save resources first because decompilation can stop or fail
appendResourcesSaveTasks(executor, resOutDir);
}
if (saveSources) {
taskCount += appendSourcesSave(tasks, sourcesOutDir, barrier);
appendSourcesSave(executor, sourcesOutDir);
}
barrier.setUpBarrier(taskCount);
return tasks;
if (gradleExportTask != null) {
executor.addSequentialTask(gradleExportTask);
}
return executor;
}
private int appendResourcesSaveTasks(List<Runnable> tasks, File outDir, TaskBarrier barrier) {
int numResourceTasks = 0;
private void appendResourcesSaveTasks(ITaskExecutor executor, File outDir) {
if (args.isSkipFilesSave()) {
return 0;
return;
}
// process AndroidManifest.xml first to load complete resource ids table
for (ResourceFile resourceFile : getResources()) {
@@ -345,8 +335,10 @@ public final class JadxDecompiler implements Closeable {
break;
}
}
Set<String> inputFileNames = args.getInputFiles().stream().map(File::getAbsolutePath).collect(Collectors.toSet());
Set<String> inputFileNames = args.getInputFiles().stream()
.map(File::getAbsolutePath)
.collect(Collectors.toSet());
List<Runnable> tasks = new ArrayList<>();
for (ResourceFile resourceFile : getResources()) {
ResourceType resType = resourceFile.getType();
if (resType == ResourceType.MANIFEST) {
@@ -358,17 +350,41 @@ public final class JadxDecompiler implements Closeable {
// ignore resource made from input file
continue;
}
tasks.add(new ResourcesSaver(outDir, resourceFile, barrier));
numResourceTasks++;
tasks.add(new ResourcesSaver(outDir, resourceFile));
}
return numResourceTasks;
executor.addParallelTasks(tasks);
}
private int appendSourcesSave(List<Runnable> tasks, File outDir, TaskBarrier barrier) {
int numSourceTasks = 0;
Predicate<String> classFilter = args.getClassFilter();
private void appendSourcesSave(ITaskExecutor executor, File outDir) {
List<JavaClass> classes = getClasses();
List<JavaClass> processQueue = new ArrayList<>(classes.size());
List<JavaClass> processQueue = filterClasses(classes);
List<List<JavaClass>> batches;
try {
batches = decompileScheduler.buildBatches(processQueue);
} catch (Exception e) {
throw new JadxRuntimeException("Decompilation batches build failed", e);
}
List<Runnable> decompileTasks = new ArrayList<>(batches.size());
for (List<JavaClass> decompileBatch : batches) {
decompileTasks.add(() -> {
for (JavaClass cls : decompileBatch) {
try {
ClassNode clsNode = cls.getClassNode();
ICodeInfo code = clsNode.getCode();
SaveCode.save(outDir, clsNode, code);
} catch (Exception e) {
LOG.error("Error saving class: {}", cls, e);
}
}
});
}
executor.addParallelTasks(decompileTasks);
}
private List<JavaClass> filterClasses(List<JavaClass> classes) {
Predicate<String> classFilter = args.getClassFilter();
List<JavaClass> list = new ArrayList<>(classes.size());
for (JavaClass cls : classes) {
ClassNode clsNode = cls.getClassNode();
if (clsNode.contains(AFlag.DONT_GENERATE)) {
@@ -380,35 +396,9 @@ public final class JadxDecompiler implements Closeable {
}
continue;
}
processQueue.add(cls);
list.add(cls);
}
List<List<JavaClass>> batches;
try {
batches = decompileScheduler.buildBatches(processQueue);
} catch (Exception e) {
throw new JadxRuntimeException("Decompilation batches build failed", e);
}
for (List<JavaClass> decompileBatch : batches) {
tasks.add(() -> {
try {
for (JavaClass cls : decompileBatch) {
try {
ClassNode clsNode = cls.getClassNode();
ICodeInfo code = clsNode.getCode();
SaveCode.save(outDir, clsNode, code);
} catch (Exception e) {
LOG.error("Error saving class: {}", cls, e);
}
}
} finally {
if (barrier != null) {
barrier.finishTask();
}
}
});
numSourceTasks++;
}
return numSourceTasks;
return list;
}
public List<JavaClass> getClasses() {
jadx-core/src/main/java/jadx/api/TaskBarrier.java
-22
View File
@@ -1,22 +0,0 @@
package jadx.api;
import java.util.concurrent.CountDownLatch;
public class TaskBarrier {
private CountDownLatch taskCountDown = null;
public void setUpBarrier(final int numTasks) {
taskCountDown = new CountDownLatch(numTasks);
}
public CountDownLatch getTaskCountDown() {
return taskCountDown;
}
public void finishTask() {
if (taskCountDown != null) {
taskCountDown.countDown();
}
}
}
jadx-core/src/main/java/jadx/api/utils/tasks/ITaskExecutor.java
+69
View File
@@ -0,0 +1,69 @@
package jadx.api.utils.tasks;
import java.util.List;
import java.util.concurrent.ExecutorService;
import org.jetbrains.annotations.Nullable;
/**
* Schedule and execute tasks combined into stages
* with parallel or sequential execution (similar to the fork-join pattern).
*/
public interface ITaskExecutor {
/**
* Add parallel stage with provided tasks
*/
void addParallelTasks(List<? extends Runnable> parallelTasks);
/**
* Add sequential stage with provided tasks
*/
void addSequentialTasks(List<? extends Runnable> seqTasks);
/**
* Add sequential stage with a single task
*/
void addSequentialTask(Runnable task);
/**
* Scheduled tasks count
*/
int getTasksCount();
/**
* Set threads count for parallel stage.
* Can be changed during execution.
* Defaults to half of processors count.
*/
void setThreadsCount(int threadsCount);
int getThreadsCount();
/**
* Start tasks execution.
*/
void execute();
int getProgress();
/**
* Not started tasks will be not executed after this method invocation.
*/
void terminate();
boolean isTerminating();
boolean isRunning();
/**
* Block until execution is finished
*/
void awaitTermination();
/**
* Return internal executor service.
*/
@Nullable
ExecutorService getInternalExecutor();
}
jadx-core/src/main/java/jadx/core/export/ExportGradleTask.java
+7 -22
View File
@@ -5,10 +5,8 @@ import java.util.List;
import jadx.api.ResourceFile;
import jadx.api.ResourceType;
import jadx.api.TaskBarrier;
import jadx.core.dex.nodes.RootNode;
import jadx.core.utils.android.AndroidManifestParser;
import jadx.core.utils.exceptions.JadxRuntimeException;
import jadx.core.utils.files.FileUtils;
import jadx.core.xmlgen.ResContainer;
@@ -18,20 +16,21 @@ public class ExportGradleTask implements Runnable {
private final RootNode root;
private final File projectDir;
private final File appDir;
private final File srcOutDir;
private final File resOutDir;
private final TaskBarrier barrier;
public ExportGradleTask(List<ResourceFile> resources, RootNode root, File projectDir, TaskBarrier barrier) {
public ExportGradleTask(List<ResourceFile> resources, RootNode root, File projectDir) {
this.resources = resources;
this.projectDir = projectDir;
this.root = root;
this.appDir = new File(projectDir, "app");
File appDir = new File(projectDir, "app");
this.srcOutDir = new File(appDir, "src/main/java");
this.resOutDir = new File(appDir, "src/main");
this.barrier = barrier;
}
public void init() {
FileUtils.makeDirs(srcOutDir);
FileUtils.makeDirs(resOutDir);
}
@Override
@@ -58,23 +57,9 @@ public class ExportGradleTask implements Runnable {
.orElse(null));
ExportGradleProject export = new ExportGradleProject(root, projectDir, androidManifest, strings);
// wait until all sources and resources are exported and all necessary info for gradle export are
// collected
try {
barrier.getTaskCountDown().await();
} catch (InterruptedException e) {
throw new JadxRuntimeException("Gradle export failed", e);
}
export.generateGradleFiles();
}
public void init() {
FileUtils.makeDirs(srcOutDir);
FileUtils.makeDirs(resOutDir);
}
public File getSrcOutDir() {
return srcOutDir;
}
jadx-core/src/main/java/jadx/core/utils/tasks/TaskExecutor.java
+180
View File
@@ -0,0 +1,180 @@
package jadx.core.utils.tasks;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import org.jetbrains.annotations.Nullable;
import jadx.api.JadxArgs;
import jadx.api.utils.tasks.ITaskExecutor;
import jadx.core.utils.exceptions.JadxRuntimeException;
public class TaskExecutor implements ITaskExecutor {
private enum ExecType {
PARALLEL,
SEQUENTIAL,
}
private static final class ExecStage {
private final ExecType type;
private final List<? extends Runnable> tasks;
private ExecStage(ExecType type, List<? extends Runnable> tasks) {
this.type = type;
this.tasks = tasks;
}
public ExecType getType() {
return type;
}
public List<? extends Runnable> getTasks() {
return tasks;
}
}
private final List<ExecStage> stages = new ArrayList<>();
private final AtomicInteger threadsCount = new AtomicInteger(JadxArgs.DEFAULT_THREADS_COUNT);
private final AtomicInteger progress = new AtomicInteger(0);
private final AtomicBoolean running = new AtomicBoolean(false);
private final AtomicBoolean terminating = new AtomicBoolean(false);
private int tasksCount = 0;
private @Nullable ExecutorService executor;
@Override
public void addParallelTasks(List<? extends Runnable> parallelTasks) {
if (parallelTasks.isEmpty()) {
return;
}
tasksCount += parallelTasks.size();
stages.add(new ExecStage(ExecType.PARALLEL, parallelTasks));
}
@Override
public void addSequentialTasks(List<? extends Runnable> seqTasks) {
if (seqTasks.isEmpty()) {
return;
}
tasksCount += seqTasks.size();
stages.add(new ExecStage(ExecType.SEQUENTIAL, seqTasks));
}
@Override
public void addSequentialTask(Runnable seqTask) {
addSequentialTasks(Collections.singletonList(seqTask));
}
@Override
public int getThreadsCount() {
return threadsCount.get();
}
@Override
public void setThreadsCount(int count) {
threadsCount.set(count);
}
@Override
public int getTasksCount() {
return tasksCount;
}
@Override
public int getProgress() {
return progress.get();
}
@Override
public void execute() {
if (running.get() || executor != null) {
throw new IllegalStateException("Already executing");
}
running.set(true);
progress.set(0);
terminating.set(false);
executor = Executors.newFixedThreadPool(1);
executor.execute(this::runStages);
executor.shutdown();
}
@Override
public void terminate() {
terminating.set(true);
}
@Override
public boolean isTerminating() {
return terminating.get();
}
@Override
public boolean isRunning() {
return running.get();
}
@Override
public @Nullable ExecutorService getInternalExecutor() {
return executor;
}
@Override
public void awaitTermination() {
if (executor == null || !running.get()) {
// already terminated
return;
}
awaitExecutorTermination(executor);
}
private void runStages() {
try {
for (ExecStage stage : stages) {
int threads = Math.min(stage.getTasks().size(), threadsCount.get());
if (stage.getType() == ExecType.SEQUENTIAL || threads == 1) {
for (Runnable task : stage.getTasks()) {
wrapTask(task);
}
} else {
ExecutorService parallelExecutor = Executors.newFixedThreadPool(threads);
for (Runnable task : stage.getTasks()) {
parallelExecutor.execute(() -> wrapTask(task));
}
parallelExecutor.shutdown();
awaitExecutorTermination(parallelExecutor);
}
if (terminating.get()) {
break;
}
}
} finally {
running.set(false);
executor = null;
}
}
private void wrapTask(Runnable task) {
if (terminating.get()) {
return;
}
task.run();
progress.incrementAndGet();
}
public static void awaitExecutorTermination(ExecutorService executor) {
try {
boolean complete = executor.awaitTermination(10, TimeUnit.DAYS);
if (!complete) {
throw new JadxRuntimeException("Executor timeout");
}
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
}
}
}
jadx-core/src/main/java/jadx/core/xmlgen/ResourcesSaver.java
-13
View File
@@ -10,7 +10,6 @@ import org.slf4j.LoggerFactory;
import jadx.api.ResourceFile;
import jadx.api.ResourcesLoader;
import jadx.api.TaskBarrier;
import jadx.api.plugins.utils.ZipSecurity;
import jadx.core.dex.visitors.SaveCode;
import jadx.core.utils.exceptions.JadxException;
@@ -23,29 +22,17 @@ public class ResourcesSaver implements Runnable {
private final ResourceFile resourceFile;
private final File outDir;
private TaskBarrier barrier = null;
public ResourcesSaver(File outDir, ResourceFile resourceFile) {
this.resourceFile = resourceFile;
this.outDir = outDir;
}
public ResourcesSaver(File outDir, ResourceFile resourceFile, TaskBarrier barrier) {
this.resourceFile = resourceFile;
this.outDir = outDir;
this.barrier = barrier;
}
@Override
public void run() {
try {
saveResources(resourceFile.loadContent());
} catch (Throwable e) {
LOG.warn("Failed to save resource: {}", resourceFile.getOriginalName(), e);
} finally {
if (barrier != null) {
barrier.finishTask();
}
}
}
jadx-core/src/test/java/jadx/tests/api/ExportGradleTest.java
+5 -7
View File
@@ -13,7 +13,6 @@ import jadx.api.ICodeInfo;
import jadx.api.JadxDecompiler;
import jadx.api.JadxDecompilerTestUtils;
import jadx.api.ResourceFile;
import jadx.api.TaskBarrier;
import jadx.core.dex.nodes.RootNode;
import jadx.core.export.ExportGradleProject;
import jadx.core.export.ExportGradleTask;
@@ -66,12 +65,11 @@ public abstract class ExportGradleTest {
final ResContainer androidManifestContainer = createResourceContainer(manifestFilename);
when(androidManifest.loadContent()).thenReturn(androidManifestContainer);
final ResContainer strings = createResourceContainer(stringsFileName);
TaskBarrier taskBarrier = new TaskBarrier();
final ExportGradleTask exportGradleTask = new ExportGradleTask(List.of(androidManifest), root, exportDir, taskBarrier);
final ExportGradleTask exportGradleTask = new ExportGradleTask(List.of(androidManifest), root, exportDir);
exportGradleTask.init();
assertThat(exportGradleTask.getSrcOutDir().exists());
assertThat(exportGradleTask.getResOutDir().exists());
assertThat(exportGradleTask.getSrcOutDir()).exists();
assertThat(exportGradleTask.getResOutDir()).exists();
final ExportGradleProject export =
new ExportGradleProject(root, exportDir, androidManifest, strings);
@@ -80,13 +78,13 @@ public abstract class ExportGradleTest {
protected String getAppGradleBuild() {
File appBuildGradle = new File(exportDir, "app/build.gradle");
assertThat(appBuildGradle.exists());
assertThat(appBuildGradle).exists();
return loadFileContent(appBuildGradle);
}
protected String getSettingsGradle() {
File settingsGradle = new File(exportDir, "settings.gradle");
assertThat(settingsGradle.exists());
assertThat(settingsGradle).exists();
return loadFileContent(settingsGradle);
}
}