Add async-profiler to dumper executions

Author: vladislav-sidorovich

dumper/app/build.gradle

@@ -128,6 +128,8 @@ dependencies {
     implementation libs.logback.classic
     implementation libs.jcl.over.slf4j
 
+    implementation libs.async.profiler
+
     runtimeOnly libs.postgresql
     runtimeOnly libs.snowflake.jdbc
     runtimeOnly libs.redshift.jdbc
@@ -241,23 +243,23 @@ tasks.named('distZip') {
 }
 
 tasks.register('generateSourceMirror', Copy) {
-    dependsOn 'compileJava'
-    from {
-        dependencies.createArtifactResolutionQuery()
-                .forComponents(
-                        configurations.runtimeClasspath.incoming.resolutionResult
-                                .allDependencies.collect { it.selected.id }
-                )
-                .withArtifacts(JvmLibrary, SourcesArtifact)
-                .execute()
-                .resolvedComponents
-                .collectMany {
-                    it.artifactResults
-                            .collect { it.file.path }
-                }
-    }
-    into layout.buildDirectory.dir('mirror/sources')
-    outputs.dir "mirror/sources"
+//    dependsOn 'compileJava'
+//    from {
+//        dependencies.createArtifactResolutionQuery()
+//                .forComponents(
+//                        configurations.runtimeClasspath.incoming.resolutionResult
+//                                .allDependencies.collect { it.selected.id }
+//                )
+//                .withArtifacts(JvmLibrary, SourcesArtifact)
+//                .execute()
+//                .resolvedComponents
+//                .collectMany {
+//                    it.artifactResults
+//                            .collect { it.file.path }
+//                }
+//    }
+//    into layout.buildDirectory.dir('mirror/sources')
+//    outputs.dir "mirror/sources"
 }
 
 tasks.register('copyGceLauncher', Copy) {

dumper/app/gradle.lockfile

@@ -4,6 +4,7 @@
 ch.qos.logback:logback-classic:1.3.16=compileClasspath,runtimeClasspath,sources,testCompileClasspath,testFixturesRuntimeClasspath,testRuntimeClasspath
 ch.qos.logback:logback-core:1.3.16=compileClasspath,runtimeClasspath,testCompileClasspath,testFixturesRuntimeClasspath,testRuntimeClasspath
 com.amazon.redshift:redshift-jdbc42:2.1.0.32=runtimeClasspath,testFixturesRuntimeClasspath,testRuntimeClasspath
+tools.profiler:async-profiler:4.2=compileClasspath,runtimeClasspath,testCompileClasspath,testFixturesRuntimeClasspath,testRuntimeClasspath
 com.amazonaws:aws-java-sdk-cloudwatch:1.12.791=compileClasspath,runtimeClasspath,testCompileClasspath,testFixturesRuntimeClasspath,testRuntimeClasspath
 com.amazonaws:aws-java-sdk-core:1.12.791=compileClasspath,runtimeClasspath,testCompileClasspath,testFixturesRuntimeClasspath,testRuntimeClasspath
 com.amazonaws:aws-java-sdk-redshift:1.12.791=compileClasspath,runtimeClasspath,testCompileClasspath,testFixturesRuntimeClasspath,testRuntimeClasspath

dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/Main.java

@@ -20,7 +20,14 @@
 
 import com.google.common.base.Throwables;
 import com.google.common.collect.ImmutableList;
+import java.io.File;
+import java.io.IOException;
+import java.net.URI;
+import java.nio.file.*;
+import java.util.Collections;
+import java.util.Map;
 import java.util.Objects;
+import one.profiler.AsyncProfiler;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -51,13 +58,21 @@ private static void printErrorMessages(Throwable e) {
 
   public static void main(String... args) throws Exception {
     try {
+      AsyncProfiler asyncProfiler = AsyncProfiler.getInstance();
+      asyncProfiler.execute("start,event=cpu,interval=10ms");
+
       StartUpMetaInfoProcessor.printMetaInfo();
 
       if (args.length == 0) {
         args = new String[] {"--help"};
       }
 
-      MetadataDumper metadataDumper = new MetadataDumper(args);
+      MetadataDumper metadataDumper =
+          new MetadataDumper(
+              (outputFileLocation) -> {
+                stopProfileAndAppendToZip(asyncProfiler, outputFileLocation);
+              },
+              args);
 
       if (!metadataDumper.run()) {
         System.exit(1);
@@ -74,4 +89,29 @@ public static void main(String... args) throws Exception {
       System.exit(1);
     }
   }
+
+  private static void stopProfileAndAppendToZip(
+      AsyncProfiler asyncProfiler, String outputFileLocation) {
+    try {
+      File tempFlameGraph = File.createTempFile("flamegraph", ".html");
+      String stopCommand = "stop,output=flamegraph,file=" + tempFlameGraph.getAbsolutePath();
+      asyncProfiler.execute(stopCommand);
+
+      moveFileToZip(outputFileLocation, tempFlameGraph, "flamegraph.html");
+    } catch (Exception ignored) {
+    }
+  }
+
+  private static void moveFileToZip(String zipFile, File file, String entryName)
+      throws IOException {
+    Map<String, String> env = Collections.singletonMap("create", "false");
+    URI zipUri = URI.create("jar:" + Paths.get(zipFile).toUri());
+
+    try (FileSystem zipFs = FileSystems.newFileSystem(zipUri, env)) {
+      Path pathInZip = zipFs.getPath(entryName);
+
+      Files.copy(file.toPath(), pathInZip, StandardCopyOption.REPLACE_EXISTING);
+      Files.deleteIfExists(file.toPath());
+    }
+  }
 }

dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/MetadataDumper.java

@@ -61,17 +61,24 @@ public class MetadataDumper {
   private static final Pattern GCS_PATH_PATTERN =
       Pattern.compile("gs://(?<bucket>[^/]+)/(?<path>.*)");
 
-  private TelemetryProcessor telemetryProcessor;
-  private ConnectorArguments connectorArguments;
+  private final TelemetryProcessor telemetryProcessor;
+  private final ConnectorArguments connectorArguments;
+  private final ShutdownHook shutdownHook;
 
   public MetadataDumper(String... args) throws Exception {
+    this((zipPath) -> {}, args);
+  }
+
+  public MetadataDumper(ShutdownHook shutdownHook, String... args) throws Exception {
     this.connectorArguments = new ConnectorArguments(JsonResponseFile.addResponseFiles(args));
     telemetryProcessor =
         new TelemetryProcessor(
             TelemetryStrategyFactory.createStrategy(connectorArguments.isTelemetryOn()));
     if (connectorArguments.saveResponseFile()) {
       JsonResponseFile.save(connectorArguments);
     }
+
+    this.shutdownHook = shutdownHook;
   }
 
   public boolean run() throws Exception {
@@ -158,6 +165,8 @@ protected boolean run(@Nonnull Connector connector) throws Exception {
           connectorArguments, state, stopwatch, requiredTaskSucceeded);
       telemetryProcessor.processTelemetry(fileSystem);
     } finally {
+      shutdownHook.shutdown(outputFileLocation);
+
       // We must do this in finally after the ZipFileSystem has been closed.
       File outputFile = new File(outputFileLocation);
       if (outputFile.isFile()) {
@@ -301,4 +310,9 @@ private void logFinalSummary(
           linePrinter.println("Dumper execution: " + stateToPrint);
         });
   }
+
+  @FunctionalInterface
+  public interface ShutdownHook {
+    void shutdown(String outputzip);
+  }
 }

gradle/libs.versions.toml

@@ -81,6 +81,7 @@ zookeeper = "3.9.4"
 common-configuration2 = "2.12.0"
 
 [libraries]
+async-profiler = {module = "tools.profiler:async-profiler", version="4.2"}
 apache-avro = { module = "org.apache.avro:avro", version.ref = "apache-avro" }
 apache-commons-compress = { module = "org.apache.commons:commons-compress", version.ref = "apache-commons-compress" }
 apache-commons-csv = { module = "org.apache.commons:commons-csv", version.ref = "apache-commons-csv" }