From 5b26876ac743a718ae7609ecd65c4514551a7f2e Mon Sep 17 00:00:00 2001
From: BodoBolero <peterbendel@neon.tech>
Date: Mon, 13 Jan 2025 19:00:04 +0100
Subject: [PATCH] test out different stripe size for ingest

---
 .../actions/neon-project-create/action.yml    | 35 +++++++++++++++++++
 .github/workflows/ingest_benchmark.yml        |  3 ++
 2 files changed, 38 insertions(+)

diff --git a/.github/actions/neon-project-create/action.yml b/.github/actions/neon-project-create/action.yml
index f4a194639f..d89ff4b31d 100644
--- a/.github/actions/neon-project-create/action.yml
+++ b/.github/actions/neon-project-create/action.yml
@@ -17,6 +17,19 @@ inputs:
   compute_units:
     description: '[Min, Max] compute units'
     default: '[1, 1]'
+  shard-split-project:
+    description: 'by default new projects are not shard-split, specify true to shard-split'
+    required: false
+    default: 'false'
+  shard-count:
+    description: 'Number of shards to split the project into, only applies if shard-split-project is true'
+    required: false
+    default: '8'
+  stripe-size:
+    description: 'Stripe size, optional, in 8kiB pages.  e.g. set 2048 for 16MB stripes. Default is 128 MiB, only applies if shard-split-project is true'
+    required: false
+    default: '32768'
+  
 
 outputs:
   dsn:
@@ -63,6 +76,24 @@ runs:
         echo "project_id=${project_id}" >> $GITHUB_OUTPUT
 
         echo "Project ${project_id} has been created"
+
+        if [ "${SHARD_SPLIT_PROJECT}" = "true" ]; then
+          # determine tenant ID
+          TENANT_ID=`${PSQL} ${dsn} -t -A -c "SHOW neon.tenant_id"`
+          
+          echo "Splitting project ${project_id} with tenant_id ${TENANT_ID} into $((SHARD_COUNT)) shards with stripe size $((STRIPE_SIZE))"
+         
+          curl -X PUT \
+            "https://${API_HOST}/regions/${REGION_ID}/api/v1/admin/storage/proxy/control/v1/tenant/${TENANT_ID}/shard_split" \
+            --header "Accept: application/json" \
+            --header "Content-Type: application/json" \
+            --header "Authorization: Bearer ${API_KEY}" \
+            --data "{
+              \"new_shard_count\": $((SHARD_COUNT)),
+              \"new_stripe_size\": $((STRIPE_SIZE))
+            }"
+        fi
+
       env:
         API_HOST: ${{ inputs.api_host }}
         API_KEY: ${{ inputs.api_key }}
@@ -70,3 +101,7 @@ runs:
         POSTGRES_VERSION: ${{ inputs.postgres_version }}
         MIN_CU: ${{ fromJSON(inputs.compute_units)[0] }}
         MAX_CU: ${{ fromJSON(inputs.compute_units)[1] }}
+        SHARD_SPLIT_PROJECT: ${{ inputs.shard-split-project }}
+        SHARD_COUNT: ${{ inputs.shard-count }}
+        STRIPE_SIZE: ${{ inputs.stripe-size }}
+        PSQL: /tmp/neon/pg_install/v16/bin/psql
diff --git a/.github/workflows/ingest_benchmark.yml b/.github/workflows/ingest_benchmark.yml
index fc33c0a980..52551d2e8c 100644
--- a/.github/workflows/ingest_benchmark.yml
+++ b/.github/workflows/ingest_benchmark.yml
@@ -75,6 +75,9 @@ jobs:
         postgres_version: 16
         compute_units: '[7, 7]' # we want to test large compute here to avoid compute-side bottleneck
         api_key: ${{ secrets.NEON_STAGING_API_KEY }}
+        shard-split-project: true
+        shard-count: 8
+        stripe-size: 512
 
     - name: Initialize Neon project
       if: ${{ matrix.target_project == 'new_empty_project' }}