Shutdown hook with heap dump added as test-only CLI option

modAionImpl/src/org/aion/utils/HeapDumper.java

@@ -24,7 +24,7 @@ public class HeapDumper {
      * @param fileName name of the heap dump file
      * @param live flag that tells whether to dump only the live objects
      */
-    static void dumpHeap(String fileName, boolean live) {
+    public static void dumpHeap(String fileName, boolean live) {
         // initialize hotspot diagnostic MBean
         initHotspotMBean();
         try {

modAionImpl/src/org/aion/zero/impl/AionBlockchainImpl.java

@@ -12,6 +12,7 @@
 import static org.aion.util.biginteger.BIUtil.isMoreThan;
 import static org.aion.util.conversions.Hex.toHexString;
 
+import java.io.File;
 import java.math.BigInteger;
 import java.util.ArrayDeque;
 import java.util.ArrayList;
@@ -63,6 +64,7 @@
 import org.aion.util.types.AddressUtils;
 import org.aion.util.types.ByteArrayWrapper;
 import org.aion.util.types.Hash256;
+import org.aion.utils.HeapDumper;
 import org.aion.vm.BlockCachingContext;
 import org.aion.vm.BulkExecutor;
 import org.aion.vm.PostExecutionLogic;
@@ -747,7 +749,20 @@ public Pair<ByteArrayWrapper, Long> findMissingAncestor(Block block) {
         }
     }
 
+    public static long shutdownHook = Long.MAX_VALUE;
+
     public synchronized ImportResult tryToConnect(final Block block) {
+        if (bestBlock.getNumber() == shutdownHook) {
+            LOG.info("Shutting down and dumping heap as indicated by CLI request since block number {} was reached.", shutdownHook);
+
+            try {
+                HeapDumper.dumpHeap(new File(System.currentTimeMillis() + "-heap-report.hprof").getAbsolutePath(), true);
+            } catch (Exception e) {
+                LOG.error("Unable to dump heap due to exception:", e);
+            }
+
+            System.exit(0);
+        }
         return tryToConnectInternal(block, System.currentTimeMillis() / THOUSAND_MS);
     }
 

modAionImpl/src/org/aion/zero/impl/cli/Arguments.java

@@ -189,6 +189,13 @@ public class Arguments {
             description = "retrieve account information")
     private String accountDetails = null;
 
+    @Option(
+            names = {"xs", "--xx-stop-at"},
+            arity = "1",
+            paramLabel = "<block_number>",
+            description = "ONLY FOR TESTS: dumps the heap and shuts down after the specified block is imported")
+    private long stopAt = Long.MAX_VALUE;
+
     /** Compacts the account options into specific commands. */
     public static String[] preProcess(String[] arguments) {
         List<String> list = new ArrayList<>();
@@ -316,4 +323,8 @@ public String getTransactionDetails() {
     public String getAccountDetails() {
         return accountDetails;
     }
+
+    public long getStopAt() {
+        return stopAt;
+    }
 }

modAionImpl/src/org/aion/zero/impl/cli/Cli.java

@@ -29,6 +29,7 @@
 import org.aion.util.bytes.ByteUtil;
 import org.aion.util.conversions.Hex;
 import org.aion.vm.LongLivedAvm;
+import org.aion.zero.impl.AionBlockchainImpl;
 import org.aion.zero.impl.Version;
 import org.aion.zero.impl.config.Network;
 import org.aion.zero.impl.db.DBUtils;
@@ -610,6 +611,15 @@ public ReturnType call(final String[] args, Cfg cfg) {
                 }
             }
 
+            if (options.getStopAt() != Long.MAX_VALUE) {
+                if (options.getStopAt() <= 0){
+                    System.out.println("Invalid parameter value for --xx-stop-at.");
+                    return ERROR;
+                }
+                System.out.println("Shutdown hook set at block " + options.getStopAt());
+                AionBlockchainImpl.shutdownHook = options.getStopAt();
+            }
+
             // if no return happened earlier, run the kernel
             return RUN;
         } catch (Exception e) {
@@ -689,6 +699,7 @@ private void printHelp() {
                 usage.replaceFirst(
                         "<slow_import>([\\s\\S]*?)]+", "| <slow_import> <frequency>\u001B[0m");
         System.out.println(usage);
+        System.out.println("Note that options prefixed with --xx are used only for testing purposes.");
     }
 
     private void printInfo(Cfg cfg) {

modAionImpl/src/org/aion/zero/impl/db/DBUtils.java

@@ -534,7 +534,6 @@ public static void redoMainChainImport(long startHeight) {
         CfgAion cfg = CfgAion.inst();
         cfg.dbFromXML();
         cfg.getConsensus().setMining(false);
-        cfg.getDb().setHeapCacheEnabled(true);
 
         System.out.println("\nImporting stored blocks INITIATED...\n");
 
@@ -617,9 +616,14 @@ public static void redoMainChainImport(long startHeight) {
                     try {
                         // clear the index entry and prune side-chain blocks
                         store.redoIndexWithoutSideChains(block);
+                        long t1 = System.currentTimeMillis();
                         result =
                                 chain.tryToConnectAndFetchSummary(
                                         block, System.currentTimeMillis() / THOUSAND_MS, false);
+                        long t2 = System.currentTimeMillis();
+                        System.out.println("<import-status: hash = " + block.getShortHash() + ", number = " + block.getNumber()
+                                               + ", txs = " + block.getTransactionsList().size() + ", result = " + result.getLeft()
+                                               + ", time elapsed = " + (t2 - t1) + " ms, td = " + chain.getTotalDifficulty() + ">");
                     } catch (Throwable t) {
                         // we want to see the exception and the block where it occurred
                         t.printStackTrace();
@@ -692,6 +696,7 @@ public AionBlockSummary getRight() {
 
                     currentBlock++;
                 }
+                System.out.println("Import from " + startHeight + " to " + topBlockNumber + " completed in " + (System.currentTimeMillis() - start) + " ms time.");
             }
 
             if (fail) {

script/prepack.sh

@@ -26,7 +26,7 @@ if [ ! -d "$JDK_PATH" ]; then
 fi
 
 module_path=$JDK_PATH/jmods 
-add_modules="java.base,java.xml,java.logging,java.management,jdk.unsupported,jdk.sctp,java.security.sasl"
+add_modules="java.base,java.xml,java.logging,java.management,jdk.unsupported,jdk.sctp,java.security.sasl,jdk.management"
 # generate aion runtime
 if [ "$useGui" = "true" ]; then
     module_path="$module_path:$JAVAFX_PATH"