From e66e3adc368801a291f3d09c778b01ce8e253187 Mon Sep 17 00:00:00 2001
From: Noel Kwan <47273164+kwannoel@users.noreply.github.com>
Date: Wed, 18 Oct 2023 01:20:15 +0800
Subject: [PATCH] fix(ci): stabilize microbenchmark (#12900)

---
 ci/scripts/run-micro-benchmarks.sh | 17 +++++++++++++++
 ci/workflows/pull-request.yml      | 33 +++++++++++++++++++++++++++++-
 2 files changed, 49 insertions(+), 1 deletion(-)

diff --git a/ci/scripts/run-micro-benchmarks.sh b/ci/scripts/run-micro-benchmarks.sh
index 13648d1661ab..568c90de425c 100755
--- a/ci/scripts/run-micro-benchmarks.sh
+++ b/ci/scripts/run-micro-benchmarks.sh
@@ -10,6 +10,12 @@ set -euo pipefail
 # Make sure the added benchmark has a unique name.
 BENCHMARKS="stream_hash_agg json_parser bench_block_iter bench_compactor bench_lru_cache bench_merge_iter"
 
+# Reference: https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instancedata-data-retrieval.html
+get_instance_type() {
+  TOKEN=`curl -X PUT "http://169.254.169.254/latest/api/token" -H "X-aws-ec2-metadata-token-ttl-seconds: 21600"` \
+  && curl -H "X-aws-ec2-metadata-token: $TOKEN" -v http://169.254.169.254/latest/meta-data/instance-type
+}
+
 # cargo criterion --bench stream_hash_agg --message-format=json
 bench() {
   BENCHMARK_NAME=$1
@@ -34,6 +40,17 @@ bench() {
 }
 
 main() {
+  # FIXME(kwannoel): This is a workaround
+  # Microbenchmarks need to be namespaced by instance types,
+  # the result upload endpoint needs to be parameterized by instance type as well to support this.
+  echo "--- Getting aws instance type"
+  local instance_type=$(get_instance_type)
+  echo "instance_type: $instance_type"
+  if [[ $instance_type != "m6i.4xlarge" ]]; then
+    echo "Only m6i.4xlarge is supported, skipping microbenchmark"
+    exit 0
+  fi
+
   # We need cargo criterion to generate machine-readable benchmark results from
   # microbench.
   echo "--- Installing cargo criterion"
diff --git a/ci/workflows/pull-request.yml b/ci/workflows/pull-request.yml
index 367aaee3730e..364455d04566 100644
--- a/ci/workflows/pull-request.yml
+++ b/ci/workflows/pull-request.yml
@@ -529,4 +529,35 @@ steps:
     if: build.pull_request.labels includes "ci/skip-ci" && !build.pull_request.draft
     commands:
       - echo "ci/skip-ci is only usable for draft Pull Requests"
-      - exit 1
\ No newline at end of file
+      - exit 1
+
+  - label: "micro benchmark"
+    command: "ci/scripts/run-micro-benchmarks.sh"
+    key: "run-micro-benchmarks"
+    if: build.pull_request.labels includes "ci/run-micro-benchmarks"
+    plugins:
+      - docker-compose#v4.9.0:
+          run: rw-build-env
+          config: ci/docker-compose.yml
+          mount-buildkite-agent: true
+    timeout_in_minutes: 60
+    retry: *auto-retry
+
+  - label: "upload micro-benchmark"
+    if: build.pull_request.labels includes "ci/run-upload-micro-benchmark"
+    command:
+      - "BUILDKITE_BUILD_NUMBER=$BUILDKITE_BUILD_NUMBER ci/scripts/upload-micro-bench-results.sh"
+    depends_on: "run-micro-benchmarks"
+    plugins:
+      - seek-oss/aws-sm#v2.3.1:
+          env:
+            BUILDKITE_TOKEN: buildkite_token
+            GITHUB_TOKEN: github-token
+      - docker-compose#v4.9.0:
+          run: rw-build-env
+          config: ci/docker-compose.yml
+          mount-buildkite-agent: true
+          environment:
+            - BUILDKITE_TOKEN
+            - GITHUB_TOKEN
+    timeout_in_minutes: 5