From 7be0b5675eccd22baaf1c4f35c3f9dd0b9160fdb Mon Sep 17 00:00:00 2001
From: Mario Fernandez <mariofer@redhat.com>
Date: Thu, 6 Jul 2023 13:42:53 +0200
Subject: [PATCH] Add topology spread constrains

Signed-off-by: Mario Fernandez <mariofer@redhat.com>
---
 Documentation/api.md                          |  1 +
 .../modules/prometheusrestrictedconfig.adoc   |  2 ++
 pkg/manifests/manifests.go                    |  5 +++
 pkg/manifests/manifests_test.go               | 34 +++++++++++++++++++
 pkg/manifests/types.go                        |  2 ++
 5 files changed, 44 insertions(+)

diff --git a/Documentation/api.md b/Documentation/api.md
index 5040545f5c..eacee03545 100644
--- a/Documentation/api.md
+++ b/Documentation/api.md
@@ -451,6 +451,7 @@ The `PrometheusRestrictedConfig` resource defines the settings for the Prometheu
 | retention | string | Defines the duration for which Prometheus retains data. This definition must be specified using the following regular expression pattern: `[0-9]+(ms\|s\|m\|h\|d\|w\|y)` (ms = milliseconds, s= seconds,m = minutes, h = hours, d = days, w = weeks, y = years). The default value is `15d`. |
 | retentionSize | string | Defines the maximum amount of disk space used by data blocks plus the write-ahead log (WAL). Supported values are `B`, `KB`, `KiB`, `MB`, `MiB`, `GB`, `GiB`, `TB`, `TiB`, `PB`, `PiB`, `EB`, and `EiB`. The default value is `nil`. |
 | tolerations | [][v1.Toleration](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.27/#toleration-v1-core) | Defines tolerations for the pods. |
+| topologySpreadConstraints | []v1.TopologySpreadConstraint | Defines a pod's topology spread constraints. |
 | volumeClaimTemplate | *[monv1.EmbeddedPersistentVolumeClaim](https://github.com/prometheus-operator/prometheus-operator/blob/v0.66.0/Documentation/api.md#embeddedpersistentvolumeclaim) | Defines persistent storage for Prometheus. Use this setting to configure the storage class and size of a volume. |
 
 [Back to TOC](#table-of-contents)
diff --git a/Documentation/openshiftdocs/modules/prometheusrestrictedconfig.adoc b/Documentation/openshiftdocs/modules/prometheusrestrictedconfig.adoc
index b509bc6164..f2803c9fdd 100644
--- a/Documentation/openshiftdocs/modules/prometheusrestrictedconfig.adoc
+++ b/Documentation/openshiftdocs/modules/prometheusrestrictedconfig.adoc
@@ -48,6 +48,8 @@ Appears in: link:userworkloadconfiguration.adoc[UserWorkloadConfiguration]
 
 |tolerations|[]v1.Toleration|Defines tolerations for the pods.
 
+|topologySpreadConstraints|[]v1.TopologySpreadConstraint|Defines a pod's topology spread constraints.
+
 |volumeClaimTemplate|*monv1.EmbeddedPersistentVolumeClaim|Defines persistent storage for Prometheus. Use this setting to configure the storage class and size of a volume.
 
 |===
diff --git a/pkg/manifests/manifests.go b/pkg/manifests/manifests.go
index c792cba6bf..7cbfe241eb 100644
--- a/pkg/manifests/manifests.go
+++ b/pkg/manifests/manifests.go
@@ -1705,6 +1705,11 @@ func (f *Factory) PrometheusUserWorkload(grpcTLS *v1.Secret, trustedCABundleCM *
 		p.Spec.Tolerations = f.config.UserWorkloadConfiguration.Prometheus.Tolerations
 	}
 
+	if len(f.config.UserWorkloadConfiguration.Prometheus.TopologySpreadConstraints) > 0 {
+		p.Spec.TopologySpreadConstraints =
+			f.config.UserWorkloadConfiguration.Prometheus.TopologySpreadConstraints
+	}
+
 	if f.config.UserWorkloadConfiguration.Prometheus.ExternalLabels != nil {
 		p.Spec.ExternalLabels = f.config.UserWorkloadConfiguration.Prometheus.ExternalLabels
 	}
diff --git a/pkg/manifests/manifests_test.go b/pkg/manifests/manifests_test.go
index a479511aa8..fd3d301909 100644
--- a/pkg/manifests/manifests_test.go
+++ b/pkg/manifests/manifests_test.go
@@ -1608,6 +1608,40 @@ ingress:
 	}
 }
 
+func TestPrometheusUserWorkloadConfiguration(t *testing.T) {
+	c := NewDefaultConfig()
+
+	uwc, err := NewUserConfigFromString(`prometheus:
+  topologySpreadConstraints:
+  - maxSkew: 1
+    topologyKey: type
+    whenUnsatisfiable: DoNotSchedule
+    labelSelector:
+      matchLabels:
+        foo: bar`)
+
+	c.UserWorkloadConfiguration = uwc
+	if err != nil {
+		t.Fatal(err)
+	}
+	f := NewFactory("openshift-monitoring", "openshift-user-workload-monitoring", c, defaultInfrastructureReader(), &fakeProxyReader{}, NewAssets(assetsPath), &APIServerConfig{}, &configv1.Console{})
+	p, err := f.PrometheusUserWorkload(
+		&v1.Secret{ObjectMeta: metav1.ObjectMeta{Name: "foo"}},
+		&v1.ConfigMap{ObjectMeta: metav1.ObjectMeta{Name: "foo"}},
+	)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if p.Spec.TopologySpreadConstraints[0].MaxSkew != 1 {
+		t.Fatal("Prometheus UWM spread contraints MaxSkew not configured correctly")
+	}
+
+	if p.Spec.TopologySpreadConstraints[0].WhenUnsatisfiable != "DoNotSchedule" {
+		t.Fatal("Prometheus UWM spread contraints WhenUnsatisfiable not configured correctly")
+	}
+}
+
 func TestPrometheusQueryLogFileConfig(t *testing.T) {
 	for _, tc := range []struct {
 		name             string
diff --git a/pkg/manifests/types.go b/pkg/manifests/types.go
index 356690fcd8..046e63caa1 100644
--- a/pkg/manifests/types.go
+++ b/pkg/manifests/types.go
@@ -586,6 +586,8 @@ type PrometheusRestrictedConfig struct {
 	RetentionSize string `json:"retentionSize,omitempty"`
 	// Defines tolerations for the pods.
 	Tolerations []v1.Toleration `json:"tolerations,omitempty"`
+	// Defines a pod's topology spread constraints.
+	TopologySpreadConstraints []v1.TopologySpreadConstraint `json:"topologySpreadConstraints,omitempty"`
 	// Defines persistent storage for Prometheus. Use this setting to
 	// configure the storage class and size of a volume.
 	VolumeClaimTemplate *monv1.EmbeddedPersistentVolumeClaim `json:"volumeClaimTemplate,omitempty"`