Support for any repo + remote exec

bench.js

@@ -10,6 +10,7 @@ const mutex = new Mutex();
 var shell = require('shelljs');
 
 var libCollector = require("./collector");
+const {benchRepo} = require("./src/bench-repo");
 
 function BenchContext(app, config) {
     var self = this;
@@ -22,7 +23,7 @@ function BenchContext(app, config) {
         const { stdout, stderr, code } = shell.exec(cmd, { silent: true });
         var error = false;
 
-        if (code != 0) {
+        if (code !== 0) {
             app.log(`ops.. Something went wrong (error code ${code})`);
             app.log(`stderr: ${stderr}`);
             error = true;
@@ -75,7 +76,7 @@ async function benchBranch(app, config) {
         collector = new libCollector.Collector();
 
         var benchContext = new BenchContext(app, config);
-        console.log(`Started benchmark "${benchConfig.title}."`);
+        app.log(`Started benchmark "${benchConfig.title}."`);
         shell.mkdir("git")
         shell.cd(cwd + "/git")
 
@@ -295,7 +296,8 @@ async function benchmarkRuntime(app, config) {
         } else if (config.repo == "polkadot") {
             benchConfig = PolkadotRuntimeBenchmarkConfigs[command];
         } else {
-            return errorResult(`${config.repo} repo is not supported.`)
+            app.log(`custom repo ${config.repo}`)
+            return benchRepo(app, config)
         }
 
         var extra = config.extra.split(" ").slice(1).join(" ").trim();
@@ -326,7 +328,7 @@ async function benchmarkRuntime(app, config) {
         }
 
         var benchContext = new BenchContext(app, config);
-        console.log(`Started runtime benchmark "${benchConfig.title}."`);
+        app.log(`Started runtime benchmark "${benchConfig.title}."`);
         shell.mkdir("git")
         shell.cd(cwd + "/git")
 

index.js

@@ -5,7 +5,9 @@ module.exports = app => {
   app.log(`base branch: ${process.env.BASE_BRANCH}`);
 
   app.on('issue_comment', async context => {
+
     let commentText = context.payload.comment.body;
+
     if (!commentText.startsWith("/bench")) {
       return;
     }
@@ -21,7 +23,9 @@ module.exports = app => {
 
     let pr = await context.github.pulls.get({ owner, repo, pull_number });
     const branchName = pr.data.head.ref;
+
     app.log(`branch: ${branchName}`);
+
     const issueComment = context.issue({ body: `Starting benchmark for branch: ${branchName} (vs ${process.env.BASE_BRANCH})\n\n Comment will be updated.` });
     const issue_comment = await context.github.issues.createComment(issueComment);
     const comment_id = issue_comment.data.id;
@@ -36,16 +40,24 @@ module.exports = app => {
       extra: extra,
     }
 
+    // Support to run the command on remote machine
+    if (process.env.REMOTE_HOST !== undefined) {
+      config.remote = { host: process.env.REMOTE_HOST,
+        user: process.env.REMOTE_USER}
+    }
+
     let report;
-    if (action == "runtime") {
+    if (action === "runtime") {
       report = await benchmarkRuntime(app, config)
     } else {
       report = await benchBranch(app, config)
-    };
+    }
 
     if (report.error) {
+
       app.log(`error: ${report.stderr}`)
-      if (report.step != "merge") {
+
+      if (report.step !== "merge") {
         context.github.issues.updateComment({
           owner, repo, comment_id,
           body: `Error running benchmark: **${branchName}**\n\n<details><summary>stdout</summary>${report.stderr}</details>`,
@@ -58,12 +70,11 @@ module.exports = app => {
       }
     } else {
       app.log(`report: ${report}`);
+
       context.github.issues.updateComment({
         owner, repo, comment_id,
         body: `Finished benchmark for branch: **${branchName}**\n\n${report}`,
       });
     }
-
-    return;
   })
 }

src/bench-context.js

@@ -0,0 +1,53 @@
+
+const shell = require('shelljs');
+
+function escq (cmd) {
+  const escaped = String.prototype.replace.call(cmd, /'/gm, "'\\''");
+  return `'${escaped}'`;
+}
+
+function BenchContext(app, config) {
+    let self = this;
+    self.app = app;
+    self.config = config;
+
+    self.temp_dir = process.env.BENCH_TEMP_DIR || 'git';
+
+    self.createTempDir = function (){
+        let cmd = `mkdir -p ${self.temp_dir}`
+        self.runTask(cmd, `Creating temp working dir ${self.temp_dir}`, false);
+    }
+
+    self.runTask = function(cmd, title, in_temp_dir=true) {
+        if (title) app.log(title);
+
+        let cmds = in_temp_dir ? `cd ${self.temp_dir} && ${cmd}` : `${cmd}`;
+
+        let cmdString = self.remoteWrapper(cmds);
+
+        const { stdout, stderr, code } = shell.exec(cmdString, { silent: true });
+        let error = false;
+
+        if (code !== 0) {
+            app.log(`ops.. Something went wrong (error code ${code})`);
+            app.log(`stderr: ${stderr}`);
+            error = true;
+        }
+
+        return { stdout, stderr, error };
+    }
+
+    self.remoteWrapper = function(cmd){
+        if (self.config.remote !== undefined) {
+            let { host, user} = config.remote;
+
+            let domain = `${user}@${host}`;
+            return `ssh ${domain} ${escq(cmd)}`;
+        }
+        return cmd;
+    }
+}
+
+module.exports = {
+    BenchContext: BenchContext
+}

src/bench-helpers.js

@@ -0,0 +1,33 @@
+function errorResult(stderr, step) {
+    return { error: true, step, stderr }
+}
+
+function checkAllowedCharacters(command) {
+    let banned = ["#", "&", "|", ";"];
+    for (const token of banned) {
+        if (command.includes(token)) {
+            return false;
+        }
+    }
+    return true;
+}
+
+function checkRuntimeBenchmarkCommand(command) {
+    let required = ["benchmark", "--pallet", "--extrinsic", "--execution", "--wasm-execution", "--steps", "--repeat", "--chain"];
+    let missing = [];
+    for (const flag of required) {
+        if (!command.includes(flag)) {
+            missing.push(flag);
+        }
+    }
+
+    return missing;
+}
+
+module.exports = {
+    errorResult: errorResult,
+    checkAllowedCharacters: checkAllowedCharacters,
+    checkRuntimeBenchmarkCommand: checkRuntimeBenchmarkCommand
+}
+
+

src/bench-repo.js

@@ -0,0 +1,158 @@
+const {BenchContext} = require("./bench-context");
+const {checkRuntimeBenchmarkCommand} = require("./bench-helpers");
+const {errorResult, checkAllowedCharacters} = require("./bench-helpers");
+
+const CustomRuntimeBenchmarkConfigs = {
+    "pallet": {
+        title: "Benchmark Runtime Pallet",
+        branchCommand: [
+            'cargo run --release',
+            '--features=runtime-benchmarks',
+            '--manifest-path={manifest_path}',
+            '--',
+            'benchmark',
+            '--chain=dev',
+            '--steps=50',
+            '--repeat=20',
+            '--pallet={pallet_name}',
+            '--extrinsic="*"',
+            '--execution=wasm',
+            '--wasm-execution=compiled',
+            '--heap-pages=4096',
+            '--output={bench_output}',
+            '--template={hbs_template}',
+        ].join(' '),
+    },
+}
+
+async function benchRepo(app, config){
+    let command = config.extra.split(" ")[0];
+
+    const supportedCommands = Object.keys(CustomRuntimeBenchmarkConfigs);
+
+    if (!supportedCommands.includes(command)){
+        return errorResult(`${command} is not supported command`)
+    }
+
+    let palletName = config.extra.split(" ").slice(1).join(" ").trim();
+
+    if (!checkAllowedCharacters(palletName)) {
+        return errorResult(`Not allowed to use #&|; in the command!`);
+    }
+
+    const manifestPath = process.env.MANIFEST_PATH || 'node/Cargo.toml';
+    const benchOutput = process.env.BENCH_PALLET_OUTPUT_FILE || 'weights.rs';
+    const hbsTemplate = process.env.BENCH_PALLET_HBS_TEMPLATE || '.maintain/pallet-weight-template.hbs';
+
+    let commandConfig = CustomRuntimeBenchmarkConfigs[command];
+
+    let cargoCommand = commandConfig.branchCommand;
+
+    cargoCommand = cargoCommand.replace("{manifest_path}", manifestPath);
+    cargoCommand = cargoCommand.replace("{bench_output}", benchOutput);
+    cargoCommand = cargoCommand.replace("{hbs_template}", hbsTemplate);
+    cargoCommand = cargoCommand.replace("{pallet_name}", palletName);
+
+    const missing = checkRuntimeBenchmarkCommand(cargoCommand);
+
+    if (missing.length > 0) {
+        return errorResult(`Missing required flags: ${missing.toString()}`)
+    }
+
+    config["title"] = commandConfig.title;
+
+    const benchContext = new BenchContext(app, config);
+
+    benchContext.palletName = palletName;
+    benchContext.benchOutput = benchOutput;
+
+    return runBench(cargoCommand, benchContext);
+}
+
+async function cloneAndSync(context){
+    let githubRepo = `https://github.com/${context.config.owner}/${context.config.repo}`;
+
+    var {error} = context.runTask(`git clone ${githubRepo} ${context.temp_dir}`, `Cloning git repository ${githubRepo} ...`, false);
+
+    if (error) {
+        context.app.log("Git clone failed, probably directory exists...");
+    }
+
+    var { stderr, error } = context.runTask(`git fetch`, "Doing git fetch...");
+
+   if (error) return errorResult(stderr);
+
+    // Checkout the custom branch
+    var { error } = context.runTask(`git checkout ${context.config.branch}`, `Checking out ${context.config.branch}...`);
+
+    if (error) {
+        context.app.log("Git checkout failed, probably some dirt in directory... Will continue with git reset.");
+    }
+
+    var { error, stderr } = context.runTask(`git reset --hard origin/${context.config.branch}`, `Resetting ${context.config.branch} hard...`);
+    if (error) return errorResult(stderr);
+
+    // Merge master branch
+    var { error, stderr } = context.runTask(`git merge origin/${context.config.baseBranch}`, `Merging branch ${context.config.baseBranch}`);
+
+    if (error) return errorResult(stderr, "merge");
+
+    if (context.config.pushToken) {
+        context.runTask(`git push https://${context.config.pushToken}@github.com/${context.config.owner}/${context.config.repo}.git HEAD`, `Pushing merge with pushToken.`);
+    } else {
+        context.runTask(`git push`, `Pushing merge.`);
+    }
+
+    return true;
+}
+
+async function runBench(command, context){
+    context.app.log(`Started runtime benchmark "${context.config.title}."`);
+
+    // If there is a preparation command - run it first
+    context.config.preparationCommand && context.runTask(context.config.preparationCommand, "Preparation command", false);
+
+    context.createTempDir();
+
+    let gitResult = await cloneAndSync(context);
+
+    if (gitResult.error){
+        return gitResult;
+    }
+
+    let { stdout, stderr, error } = context.runTask(command, `Benching branch: ${context.config.branch}...`);
+
+    if (error){
+        return errorResult(stderr, 'benchmark')
+    }
+
+    let output = command.includes("--output");
+
+    // If `--output` is set, we commit the benchmark file to the repo
+    if (output) {
+        let palletFolder = context.palletName.split('_').join('-').trim();
+        let weightsPath = `pallets/${palletFolder}/src/weights.rs`;
+        let cmd = `mv ${context.benchOutput} ${weightsPath}`;
+
+        context.runTask(cmd);
+
+        context.runTask(`git add ${weightsPath}`, `Adding new files.`);
+        context.runTask(`git commit -m "Weights update for ${context.palletName} pallet"`, `Committing changes.`);
+
+        if (config.pushToken) {
+            context.runTask(`git push https://${context.config.pushToken}@github.com/${context.config.owner}/${context.config.repo}.git HEAD`, `Pushing commit with pushToken.`);
+        } else {
+            context.runTask(`git push origin ${context.config.branch}`, `Pushing commit.`);
+        }
+    }
+
+    return `Benchmark: **${context.config.title}**\n\n`
+        + command
+        + "\n\n<details>\n<summary>Results</summary>\n\n"
+        + (stdout ? stdout : stderr)
+        + "\n\n </details>";
+}
+
+module.exports = {
+    benchRepo : benchRepo
+}

src/env.template

@@ -0,0 +1,19 @@
+WEBHOOK_PROXY_URL=https://smee.io/c3QrFmKxYuLn1jx
+
+# Github APP
+APP_ID=""
+PRIVATE_KEY_PATH=""
+WEBHOOK_SECRET=""
+BASE_BRANCH="main"
+
+# Cargo command options to replace with
+MANIFEST_PATH="node/Cargo.toml"
+BENCH_PALLET_OUTPUT_FILE="weights.rs"
+BENCH_PALLET_HBS_TEMPLATE=".maintain/pallet-weight-template.hbs"
+
+# Directory which repository is cloned into
+BENCH_TEMP_DIR="git"
+
+# Remote support - if specified - all commands are executed on the remote machine instead locally
+#REMOTE_HOST=""
+#REMOTE_USER=""