From bf4750d457b841ad8182aca967ecee4977fb4009 Mon Sep 17 00:00:00 2001
From: Emily Chui <echui@gitlab.com>
Date: Wed, 18 Sep 2024 21:58:24 +1000
Subject: [PATCH] Add ability to configure number of replicas in Gitaly's
 statefulset

---
 .../gitaly/templates/_statefulset_spec.yaml   |  2 +-
 doc/charts/gitlab/gitaly/index.md             |  1 +
 spec/configuration/gitaly_spec.rb             | 39 +++++++++++++++++++
 3 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/charts/gitlab/charts/gitaly/templates/_statefulset_spec.yaml b/charts/gitlab/charts/gitaly/templates/_statefulset_spec.yaml
index 9d72f30190..9293fb9511 100644
--- a/charts/gitlab/charts/gitaly/templates/_statefulset_spec.yaml
+++ b/charts/gitlab/charts/gitaly/templates/_statefulset_spec.yaml
@@ -16,7 +16,7 @@ spec:
   {{- if .storage }}
   replicas: {{ default (include "gitlab.gitaly.replicas" .) .storage.gitalyReplicas }}
   {{- else }}
-  replicas: {{ include "gitlab.gitaly.replicas" . }}
+  replicas: {{ coalesce .Values.statefulset.replicas (include "gitlab.gitaly.replicas" .) }}
   {{- end }}
   podManagementPolicy: Parallel
   {{- if .Values.statefulset.strategy }}
diff --git a/doc/charts/gitlab/gitaly/index.md b/doc/charts/gitlab/gitaly/index.md
index f9ddaa2475..1855d221f9 100644
--- a/doc/charts/gitlab/gitaly/index.md
+++ b/doc/charts/gitlab/gitaly/index.md
@@ -109,6 +109,7 @@ the `helm install` command using the `--set` flags.
 | `git.config[]`                                   | `[]`                                              | Git configuration that Gitaly should set when spawning Git commands                                                                                                            |
 | `prometheus.grpcLatencyBuckets`                  |                                                   | Buckets corresponding to histogram latencies on GRPC method calls to be recorded by Gitaly. A string form of the array (for example, `"[1.0, 1.5, 2.0]"`) is required as input |
 | `statefulset.strategy`                           | `{}`                                              | Allows one to configure the update strategy utilized by the StatefulSet                                                                                                        |
+| `statefulset.replicas`                           |                                                   | If not set, will default to the number of `gitaly.internal.names` or 0 if `global.gitaly.host` is set.                                                                         |
 | `statefulset.livenessProbe.initialDelaySeconds`  | 0                                                | Delay before liveness probe is initiated. If startupProbe is enabled, this will be set to 0.                                                                                                                                       |
 | `statefulset.livenessProbe.periodSeconds`        | 10                                                | How often to perform the liveness probe                                                                                                                                        |
 | `statefulset.livenessProbe.timeoutSeconds`       | 3                                                 | When the liveness probe times out                                                                                                                                              |
diff --git a/spec/configuration/gitaly_spec.rb b/spec/configuration/gitaly_spec.rb
index 6b0c7bebf9..6b9609152a 100644
--- a/spec/configuration/gitaly_spec.rb
+++ b/spec/configuration/gitaly_spec.rb
@@ -761,6 +761,45 @@ describe 'Gitaly configuration' do
     end
   end
 
+  context 'gitaly statefulset replicas' do
+    let(:values) do
+      YAML.safe_load(%(
+        global:
+          gitaly:
+            enabled: true
+            internal:
+              names:
+                - default
+                - gitaly-store-01
+                - gitaly-store-02
+        gitlab:
+          gitaly:
+            statefulset:
+              replicas: #{gitaly_sts_replicas}
+      )).merge(default_values)
+    end
+
+    let(:template) { HelmTemplate.new(values) }
+    let(:gitaly_stateful_set) { template['StatefulSet/test-gitaly'] }
+    let(:gitaly_replicas) { gitaly_stateful_set['spec']['replicas'] }
+
+    context 'when default' do
+      let(:gitaly_sts_replicas) {}
+
+      it 'sets the replicas to number of internal storage' do
+        expect(gitaly_replicas).to eq(3)
+      end
+    end
+
+    context 'when replica is provided' do
+      let(:gitaly_sts_replicas) { 5 }
+
+      it 'sets the statefulset.replicas' do
+        expect(gitaly_replicas).to eq(5)
+      end
+    end
+  end
+
   context 'gitaly service' do
     let(:values) do
       YAML.safe_load(%(
-- 
GitLab