feat(nix): add stage 1 AMI build caching based on input hash

Implement content-based caching for stage 1 AMI builds by computing a hash from all input sources.
Then the build-ami script checks for existing AMIs with matching input hash before building.

Author: jfroche

.github/workflows/testinfra-ami-build.yml

@@ -107,15 +107,23 @@ jobs:
           echo 'postgres-version = "'$PG_VERSION'"' > common-nix.vars.pkr.hcl
           # Ensure there's a newline at the end of the file
           echo "" >> common-nix.vars.pkr.hcl
+          git add -f common-nix.vars.pkr.hcl
 
       - name: Build AMI stage 1
         env:
           AWS_MAX_ATTEMPTS: 10
           AWS_RETRY_MODE: adaptive
         run: |
           GIT_SHA=${{github.sha}}
-          nix run github:supabase/postgres/${GIT_SHA}#packer -- init amazon-arm64-nix.pkr.hcl
-          nix run github:supabase/postgres/${GIT_SHA}#packer -- build -var "git-head-version=${GIT_SHA}" -var "packer-execution-id=${EXECUTION_ID}" -var-file="development-arm.vars.pkr.hcl" -var-file="common-nix.vars.pkr.hcl" -var "ansible_arguments=" -var "postgres-version=${{ steps.random.outputs.random_string }}" -var "region=ap-southeast-1" -var 'ami_regions=["ap-southeast-1"]' -var "force-deregister=true" -var "ansible_arguments=-e postgresql_major=${POSTGRES_MAJOR_VERSION}" amazon-arm64-nix.pkr.hcl
+          nix run .#build-ami -- \
+            -var "git-head-version=${GIT_SHA}" \
+            -var "packer-execution-id=${EXECUTION_ID}" \
+            -var "ansible_arguments=-e postgresql_major=${POSTGRES_MAJOR_VERSION}" \
+            -var "region=ap-southeast-1" \
+            -var 'ami_regions=["ap-southeast-1"]' \
+            -var-file="development-arm.vars.pkr.hcl" \
+            -var-file="common-nix.vars.pkr.hcl" \
+            amazon-arm64-nix.pkr.hcl
 
       - name: Build AMI stage 2
         env:

amazon-arm64-nix.pkr.hcl

@@ -92,6 +92,12 @@ variable "force-deregister" {
   default = false
 }
 
+variable "input-hash" {
+  type    = string
+  default = ""
+  description = "Content hash of all input sources"
+}
+
 packer {
   required_plugins {
     amazon = {
@@ -172,6 +178,7 @@ source "amazon-ebssurrogate" "source" {
     appType = "postgres"
     postgresVersion = "${var.postgres-version}-stage1"
     sourceSha = "${var.git-head-version}"
+    inputHash = "${var.input-hash}"
   }
 
   communicator = "ssh"

nix/packages/build-ami.nix

@@ -0,0 +1,101 @@
+{
+  lib,
+  stdenv,
+  writeShellApplication,
+  packer,
+  awscli2,
+  jq,
+  ...
+}:
+
+let
+  root = ../..;
+  packerSources = stdenv.mkDerivation {
+    name = "packer-sources";
+    src = lib.fileset.toSource {
+      inherit root;
+      fileset = lib.fileset.unions [
+        (root + "/ebssurrogate")
+        (root + "/ansible")
+        (root + "/migrations")
+        (root + "/scripts")
+        (root + "/amazon-arm64-nix.pkr.hcl")
+        (root + "/development-arm.vars.pkr.hcl")
+        (lib.fileset.maybeMissing (root + "/common-nix.vars.pkr.hcl"))
+      ];
+    };
+
+    phases = [
+      "unpackPhase"
+      "installPhase"
+    ];
+    installPhase = ''
+      mkdir -p $out
+      cp -r . $out/
+    '';
+  };
+in
+writeShellApplication {
+  name = "build-ami";
+
+  runtimeInputs = [
+    packer
+    awscli2
+    jq
+  ];
+
+  text = ''
+    set -euo pipefail
+
+    set -x
+    REGION="''${AWS_REGION:-ap-southeast-1}"
+    POSTGRES_VERSION="''${POSTGRES_MAJOR_VERSION:-15}"
+    PACKER_SOURCES="${packerSources}"
+    INPUT_HASH=$(basename "$PACKER_SOURCES" | cut -d- -f1)
+
+    echo "Checking for existing AMI..."
+    set +e
+    AMI_OUTPUT=$(aws ec2 describe-images \
+      --region "$REGION" \
+      --owners self \
+      --filters \
+        "Name=tag:inputHash,Values=$INPUT_HASH" \
+        "Name=tag:postgresVersion,Values=$POSTGRES_VERSION-stage1" \
+        "Name=state,Values=available" \
+      --query 'Images[0].ImageId' \
+      --output text 2>&1)
+    AWS_EXIT_CODE=$?
+    set -e
+
+    if [ $AWS_EXIT_CODE -ne 0 ] && [ $AWS_EXIT_CODE -ne 255 ]; then
+      echo "Error querying AWS: $AMI_OUTPUT"
+      exit 1
+    fi
+
+    AMI_ID="$AMI_OUTPUT"
+    if [ "$AMI_ID" != "None" ] && [ -n "$AMI_ID" ]; then
+      echo "Found existing AMI: $AMI_ID"
+      exit 0
+    fi
+
+    echo "No cached AMI found"
+
+    cd "$PACKER_SOURCES"
+    packer init amazon-arm64-nix.pkr.hcl
+    packer build \
+      -var "input-hash=$INPUT_HASH" \
+      -var "postgres-version=$POSTGRES_VERSION" \
+      -var "region=$REGION" \
+      "$@"
+  '';
+
+  meta = {
+    description = "Build AMI if not cached based on input hash";
+    longDescription = ''
+      The input hash is computed from all source files that affect the build.
+      Before building, we verify the existence of an AMI with the same hash.
+      If found, the build is skipped. Otherwise, a new AMI is created and
+      tagged with the input hash for future cache hits.
+    '';
+  };
+}

nix/packages/default.nix

@@ -29,6 +29,7 @@
     {
       packages = (
         {
+          build-ami = pkgs.callPackage ./build-ami.nix { packer = self'.packages.packer; };
           build-test-ami = pkgs.callPackage ./build-test-ami.nix { };
           cleanup-ami = pkgs.callPackage ./cleanup-ami.nix { };
           dbmate-tool = pkgs.callPackage ./dbmate-tool.nix { inherit (self.supabase) defaults; };