xataio · ryanslade · Nov 18, 2024 · Nov 12, 2024 · Nov 12, 2024 · Nov 12, 2024
diff --git a/.github/workflows/benchmark.yaml b/.github/workflows/benchmark.yaml
@@ -1,9 +1,11 @@
 name: Benchmark
 on:
+  workflow_dispatch:
   push:
     branches:
       - main
 permissions:
+  id-token: write # For getting AWS permissions
   contents: read
   packages: read
 jobs:
@@ -25,4 +27,37 @@ jobs:
     - name: Run benchmarks
       run: make bench
       env:
-        POSTGRES_VERSION: ${{ matrix.pgVersion }}
+        POSTGRES_VERSION: ${{ matrix.pgVersion }}
+
+    - name: Upload results
+      uses: actions/upload-artifact@v4
+      with:
+        name: benchmark_result_${{ matrix.pgVersion }}.json
+        path: internal/benchmarks/benchmark_result_${{ matrix.pgVersion }}.json
+
+  gather:
+    name: 'Gather results'
+    runs-on: ubuntu-latest
+    needs: [benchmark]
+
+    steps:
+    - uses: actions/download-artifact@v4
+      with:
+        path: ./results/
+        merge-multiple: true
+
+    - name: Configure AWS Credentials
+      uses: aws-actions/configure-aws-credentials@v2
+      with:
+        role-to-assume: arn:aws:iam::493985724844:role/pgroll-benchmark-results-access
+        aws-region: us-east-1
+        mask-aws-account-id: 'no'
+
+    - name: Download current results from S3
+      run: aws s3 cp s3://pgroll-benchmark-results/benchmark-results.json ./benchmark-results.json
+
+    - name: Append new results
+      run: cat results/*.json >> benchmark-results.json
+
+    - name: Upload combined results
+      run: aws s3 cp ./benchmark-results.json s3://pgroll-benchmark-results/benchmark-results.json
diff --git a/internal/benchmarks/benchmarks.go b/internal/benchmarks/benchmarks.go
@@ -0,0 +1,39 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package benchmarks
+
+import (
+	"os"
+	"time"
+)
+
+type Reports struct {
+	GitSHA          string
+	PostgresVersion string
+	Timestamp       int64
+	Reports         []Report
+}
+
+func (r *Reports) AddReport(report Report) {
+	r.Reports = append(r.Reports, report)
+}
+
+func getPostgresVersion() string {
+	return os.Getenv("POSTGRES_VERSION")
+}
+
+func newReports() *Reports {
+	return &Reports{
+		GitSHA:          os.Getenv("GITHUB_SHA"),
+		PostgresVersion: getPostgresVersion(),
+		Timestamp:       time.Now().Unix(),
+		Reports:         []Report{},
+	}
+}
+
+type Report struct {
+	Name     string
+	RowCount int
+	Unit     string
+	Result   float64
+}
diff --git a/internal/benchmarks/benchmarks_test.go b/internal/benchmarks/benchmarks_test.go
@@ -5,6 +5,9 @@ package benchmarks
 import (
 	"context"
 	"database/sql"
+	"encoding/json"
+	"fmt"
+	"os"
 	"strconv"
 	"testing"
 
@@ -19,10 +22,32 @@ import (
 
 const unitRowsPerSecond = "rows/s"
 
-var rowCounts = []int{10_000, 100_000, 300_000}
+var (
+	rowCounts = []int{10_000, 100_000, 300_000}
+	reports   = newReports()
+)
 
 func TestMain(m *testing.M) {
-	testutils.SharedTestMain(m)
+	testutils.SharedTestMain(m, func() (err error) {
+		// Only run in GitHub actions
+		if os.Getenv("GITHUB_ACTIONS") != "true" {
+			return nil
+		}
+
+		w, err := os.Create(fmt.Sprintf("benchmark_result_%s.json", getPostgresVersion()))
+		if err != nil {
+			return fmt.Errorf("creating report file: %w", err)
+		}
+		defer func() {
+			err = w.Close()
+		}()
+
+		encoder := json.NewEncoder(w)
+		if err := encoder.Encode(reports); err != nil {
+			return fmt.Errorf("encoding report file: %w", err)
+		}
+		return nil
+	})
 }
 
 func BenchmarkBackfill(b *testing.B) {
@@ -48,6 +73,8 @@ func BenchmarkBackfill(b *testing.B) {
 				b.Logf("Backfilled %d rows in %s", rowCount, b.Elapsed())
 				rowsPerSecond := float64(rowCount) / b.Elapsed().Seconds()
 				b.ReportMetric(rowsPerSecond, unitRowsPerSecond)
+
+				addRowsPerSecond(b, rowCount, rowsPerSecond)
 			})
 		})
 	}
@@ -87,6 +114,8 @@ func BenchmarkWriteAmplification(b *testing.B) {
 					b.StopTimer()
 					rowsPerSecond := float64(rowCount) / b.Elapsed().Seconds()
 					b.ReportMetric(rowsPerSecond, unitRowsPerSecond)
+
+					addRowsPerSecond(b, rowCount, rowsPerSecond)
 				})
 			})
 		}
@@ -116,6 +145,8 @@ func BenchmarkWriteAmplification(b *testing.B) {
 					b.StopTimer()
 					rowsPerSecond := float64(rowCount) / b.Elapsed().Seconds()
 					b.ReportMetric(rowsPerSecond, unitRowsPerSecond)
+
+					addRowsPerSecond(b, rowCount, rowsPerSecond)
 				})
 			})
 		}
@@ -149,6 +180,15 @@ func setupInitialTable(tb testing.TB, ctx context.Context, testSchema string, mi
 	seed(tb, rowCount, db)
 }
 
+func addRowsPerSecond(b *testing.B, rowCount int, perSecond float64) {
+	reports.AddReport(Report{
+		Name:     b.Name(),
+		Unit:     unitRowsPerSecond,
+		RowCount: rowCount,
+		Result:   perSecond,
+	})
+}
+
 // Simple table with a nullable `name` field.
 var migCreateTable = migrations.Migration{
 	Name: "01_create_table",

diff --git a/internal/testutils/util.go b/internal/testutils/util.go
@@ -28,9 +28,11 @@ const defaultPostgresVersion = "15.3"
 // tConnStr holds the connection string to the test container created in TestMain.
 var tConnStr string
 
-// SharedTestMain starts a postgres container to be used by all tests in a package.
-// Each test then connects to the container and creates a new database.
-func SharedTestMain(m *testing.M) {
+// SharedTestMain starts a postgres container to be used by all tests in a package. Each test then
+// connects to the container and creates a new database. Optional functions that will run after all
+// tests can be added and should return a nil error to indicate they ran successfully. If they return
+// an error all subsequent functions will be skipped.
+func SharedTestMain(m *testing.M, postRunHooks ...func() error) {
 	ctx := context.Background()
 
 	waitForLogs := wait.
@@ -73,7 +75,20 @@ func SharedTestMain(m *testing.M) {
 		log.Printf("Failed to terminate container: %v", err)
 	}
 
-	os.Exit(exitCode)
+	if exitCode != 0 {
+		log.Printf("Non zero exit code (%d), skipping post run hooks", exitCode)
+		os.Exit(exitCode)
+	}
+
+	for _, hook := range postRunHooks {
+		err := hook()
+		if err != nil {
+			log.Printf("Post-run hook failed: %v", err)
+			os.Exit(1)
+		}
+	}
+
+	os.Exit(0)
 }
 
 // TestSchema returns the schema in which migration tests apply migrations. By