Chore: add some jsonnet utilities to reduce code duplication

operations/mimir/alertmanager.libsonnet

@@ -1,7 +1,6 @@
 {
   local configMap = $.core.v1.configMap,
   local container = $.core.v1.container,
-  local podDisruptionBudget = $.policy.v1beta1.podDisruptionBudget,
   local pvc = $.core.v1.persistentVolumeClaim,
   local service = $.core.v1.service,
   local statefulSet = $.apps.v1.statefulSet,
@@ -66,17 +65,9 @@
 
   alertmanager_statefulset:
     if $._config.alertmanager_enabled then
-      statefulSet.new('alertmanager', $._config.alertmanager.replicas, [$.alertmanager_container], $.alertmanager_pvc) +
-      statefulSet.mixin.spec.withServiceName('alertmanager') +
-      statefulSet.mixin.metadata.withNamespace($._config.namespace) +
-      statefulSet.mixin.metadata.withLabels({ name: 'alertmanager' }) +
-      statefulSet.mixin.spec.template.metadata.withLabels({ name: 'alertmanager' }) +
-      statefulSet.mixin.spec.selector.withMatchLabels({ name: 'alertmanager' }) +
-      statefulSet.mixin.spec.template.spec.securityContext.withRunAsUser(0) +
-      statefulSet.mixin.spec.updateStrategy.withType('RollingUpdate') +
+      $.newMimirStatefulSet('alertmanager', $._config.alertmanager.replicas, $.alertmanager_container, $.alertmanager_pvc, podManagementPolicy=null) +
       statefulSet.mixin.spec.template.spec.withTerminationGracePeriodSeconds(900) +
       $.util.configVolumeMount($._config.overrides_configmap, $._config.overrides_configmap_mountpoint) +
-      (if !std.isObject($._config.node_selector) then {} else statefulSet.mixin.spec.template.spec.withNodeSelectorMixin($._config.node_selector)) +
       statefulSet.mixin.spec.template.spec.withVolumesMixin(
         if hasFallbackConfig then
           [volume.fromConfigMap('alertmanager-fallback-config', 'alertmanager-fallback-config')]
@@ -90,13 +81,6 @@
       service.mixin.spec.withClusterIp('None')
     else {},
 
-  alertmanager_pdb:
-    if $._config.alertmanager_enabled then
-      podDisruptionBudget.new('alertmanager-pdb') +
-      podDisruptionBudget.mixin.metadata.withLabels({ name: 'alertmanager-pdb' }) +
-      podDisruptionBudget.mixin.spec.selector.withMatchLabels({
-        name: $.alertmanager_statefulset.spec.template.metadata.labels.name,
-      }) +
-      podDisruptionBudget.mixin.spec.withMaxUnavailable(1)
-    else {},
+  alertmanager_pdb: if !$._config.alertmanager_enabled then null else
+    $.newMimirPdb('alertmanager'),
 }

operations/mimir/common.libsonnet

@@ -19,11 +19,43 @@
       container.mixin.readinessProbe.withTimeoutSeconds(1),
   },
 
-  newDiscoveryService(name, deployment)::
+  // Utility to create an headless service used to discover replicas of a Mimir deployment.
+  newMimirDiscoveryService(name, deployment)::
     local service = $.core.v1.service;
 
     $.util.serviceFor(deployment, $._config.service_ignored_labels) +
     service.mixin.spec.withPublishNotReadyAddresses(true) +
     service.mixin.spec.withClusterIp('None') +
     service.mixin.metadata.withName(name),
+
+  // Utility to create a PodDisruptionBudget for a Mimir deployment.
+  newMimirPdb(deploymentName, maxUnavailable=1)::
+    local podDisruptionBudget = $.policy.v1beta1.podDisruptionBudget;
+    local pdbName = '%s-pdb' % deploymentName;
+
+    podDisruptionBudget.new() +
+    podDisruptionBudget.mixin.metadata.withName(pdbName) +
+    podDisruptionBudget.mixin.metadata.withLabels({ name: pdbName }) +
+    podDisruptionBudget.mixin.spec.selector.withMatchLabels({ name: deploymentName }) +
+    podDisruptionBudget.mixin.spec.withMaxUnavailable(maxUnavailable),
+
+  // Utility to create a StatefulSet for a Mimir component.
+  //
+  // By default, it uses the Parallel pod management policy which parallelly
+  // scale up/down instances instead of starting them one by one. This does NOT
+  // affect rolling updates: they will continue to be rolled out one by one
+  // (the next pod will be rolled out once the previous is ready).
+  newMimirStatefulSet(name, replicas, container, pvc, podManagementPolicy='Parallel')::
+    local statefulSet = $.apps.v1.statefulSet;
+
+    statefulSet.new(name, replicas, [container], pvc) +
+    statefulSet.mixin.spec.withServiceName(name) +
+    statefulSet.mixin.metadata.withNamespace($._config.namespace) +
+    statefulSet.mixin.metadata.withLabels({ name: name }) +
+    statefulSet.mixin.spec.template.metadata.withLabels({ name: name }) +
+    statefulSet.mixin.spec.selector.withMatchLabels({ name: name }) +
+    statefulSet.mixin.spec.template.spec.securityContext.withRunAsUser(0) +
+    statefulSet.mixin.spec.updateStrategy.withType('RollingUpdate') +
+    (if podManagementPolicy != null then statefulSet.mixin.spec.withPodManagementPolicy(podManagementPolicy) else {}) +
+    (if !std.isObject($._config.node_selector) then {} else statefulSet.mixin.spec.template.spec.withNodeSelectorMixin($._config.node_selector)),
 }

operations/mimir/compactor.libsonnet

@@ -104,21 +104,8 @@
     $.jaeger_mixin,
 
   newCompactorStatefulSet(name, container)::
-    statefulSet.new(name, 1, [container], compactor_data_pvc) +
-    statefulSet.mixin.spec.withServiceName(name) +
-    statefulSet.mixin.metadata.withNamespace($._config.namespace) +
-    statefulSet.mixin.metadata.withLabels({ name: name }) +
-    statefulSet.mixin.spec.template.metadata.withLabels({ name: name }) +
-    statefulSet.mixin.spec.selector.withMatchLabels({ name: name }) +
-    statefulSet.mixin.spec.template.spec.securityContext.withRunAsUser(0) +
-    statefulSet.mixin.spec.updateStrategy.withType('RollingUpdate') +
+    $.newMimirStatefulSet(name, 1, container, compactor_data_pvc) +
     statefulSet.mixin.spec.template.spec.withTerminationGracePeriodSeconds(900) +
-    (if !std.isObject($._config.node_selector) then {} else statefulSet.mixin.spec.template.spec.withNodeSelectorMixin($._config.node_selector)) +
-    // Parallelly scale up/down compactor instances instead of starting them
-    // one by one. This does NOT affect rolling updates: they will continue to be
-    // rolled out one by one (the next pod will be rolled out once the previous is
-    // ready).
-    statefulSet.mixin.spec.withPodManagementPolicy('Parallel') +
     $.util.configVolumeMount($._config.overrides_configmap, $._config.overrides_configmap_mountpoint),
 
   compactor_statefulset:

operations/mimir/ingester.libsonnet

@@ -1,6 +1,5 @@
 {
   local container = $.core.v1.container,
-  local podDisruptionBudget = $.policy.v1beta1.podDisruptionBudget,
   local pvc = $.core.v1.persistentVolumeClaim,
   local statefulSet = $.apps.v1.statefulSet,
   local volumeMount = $.core.v1.volumeMount,
@@ -63,42 +62,25 @@
     pvc.mixin.metadata.withName('ingester-data'),
 
   newIngesterStatefulSet(name, container, with_anti_affinity=true)::
-    statefulSet.new(name, 3, [
-      container + $.core.v1.container.withVolumeMountsMixin([
-        volumeMount.new('ingester-data', '/data'),
-      ]),
-    ], ingester_data_pvc) +
-    statefulSet.mixin.spec.withServiceName(name) +
-    statefulSet.mixin.metadata.withNamespace($._config.namespace) +
-    statefulSet.mixin.metadata.withLabels({ name: name }) +
-    statefulSet.mixin.spec.template.metadata.withLabels({ name: name }) +
-    statefulSet.mixin.spec.selector.withMatchLabels({ name: name }) +
-    statefulSet.mixin.spec.template.spec.securityContext.withRunAsUser(0) +
+    local ingesterContainer = container + $.core.v1.container.withVolumeMountsMixin([
+      volumeMount.new('ingester-data', '/data'),
+    ]);
+
+    $.newMimirStatefulSet(name, 3, ingesterContainer, ingester_data_pvc) +
     // When the ingester needs to flush blocks to the storage, it may take quite a lot of time.
     // For this reason, we grant an high termination period (80 minutes).
     statefulSet.mixin.spec.template.spec.withTerminationGracePeriodSeconds(1200) +
-    statefulSet.mixin.spec.updateStrategy.withType('RollingUpdate') +
     $.util.configVolumeMount($._config.overrides_configmap, $._config.overrides_configmap_mountpoint) +
     $.util.podPriority('high') +
-    // Parallelly scale up/down ingester instances instead of starting them
-    // one by one. This does NOT affect rolling updates: they will continue to be
-    // rolled out one by one (the next pod will be rolled out once the previous is
-    // ready).
-    statefulSet.mixin.spec.withPodManagementPolicy('Parallel') +
-    (if !std.isObject($._config.node_selector) then {} else statefulSet.mixin.spec.template.spec.withNodeSelectorMixin($._config.node_selector)) +
     (if with_anti_affinity then $.util.antiAffinity else {}),
 
   ingester_statefulset: self.newIngesterStatefulSet('ingester', $.ingester_container, !$._config.ingester_allow_multiple_replicas_on_same_node),
 
   ingester_service:
     $.util.serviceFor($.ingester_statefulset, $._config.service_ignored_labels),
 
-  newIngesterPdb(pdbName, ingesterName)::
-    podDisruptionBudget.new() +
-    podDisruptionBudget.mixin.metadata.withName(pdbName) +
-    podDisruptionBudget.mixin.metadata.withLabels({ name: pdbName }) +
-    podDisruptionBudget.mixin.spec.selector.withMatchLabels({ name: ingesterName }) +
-    podDisruptionBudget.mixin.spec.withMaxUnavailable(1),
+  newIngesterPdb(ingesterName)::
+    $.newMimirPdb(ingesterName),
 
-  ingester_pdb: self.newIngesterPdb('ingester-pdb', name),
+  ingester_pdb: self.newIngesterPdb(name),
 }

operations/mimir/query-frontend.libsonnet

@@ -59,5 +59,5 @@
     // each query-frontend pod IP and NOT the service IP. To make it, we do NOT
     // use the service cluster IP so that when the service DNS is resolved it
     // returns the set of query-frontend IPs.
-    $.newDiscoveryService('query-frontend-discovery', $.query_frontend_deployment),
+    $.newMimirDiscoveryService('query-frontend-discovery', $.query_frontend_deployment),
 }

operations/mimir/query-scheduler.libsonnet

@@ -43,7 +43,7 @@
 
   // Headless to make sure resolution gets IP address of target pods, and not service IP.
   newQuerySchedulerDiscoveryService(name, deployment)::
-    $.newDiscoveryService(discoveryServiceName(name), deployment),
+    $.newMimirDiscoveryService(discoveryServiceName(name), deployment),
 
   query_scheduler_discovery_service: if !$._config.query_scheduler_enabled then {} else
     self.newQuerySchedulerDiscoveryService('query-scheduler', $.query_scheduler_deployment),

operations/mimir/store-gateway.libsonnet

@@ -1,6 +1,5 @@
 {
   local container = $.core.v1.container,
-  local podDisruptionBudget = $.policy.v1beta1.podDisruptionBudget,
   local pvc = $.core.v1.persistentVolumeClaim,
   local statefulSet = $.apps.v1.statefulSet,
   local volumeMount = $.core.v1.volumeMount,
@@ -67,21 +66,8 @@
     $.jaeger_mixin,
 
   newStoreGatewayStatefulSet(name, container, with_anti_affinity=false)::
-    statefulSet.new(name, 3, [container], store_gateway_data_pvc) +
-    statefulSet.mixin.spec.withServiceName(name) +
-    statefulSet.mixin.metadata.withNamespace($._config.namespace) +
-    statefulSet.mixin.metadata.withLabels({ name: name }) +
-    statefulSet.mixin.spec.template.metadata.withLabels({ name: name }) +
-    statefulSet.mixin.spec.selector.withMatchLabels({ name: name }) +
-    statefulSet.mixin.spec.template.spec.securityContext.withRunAsUser(0) +
-    (if !std.isObject($._config.node_selector) then {} else statefulSet.mixin.spec.template.spec.withNodeSelectorMixin($._config.node_selector)) +
-    statefulSet.mixin.spec.updateStrategy.withType('RollingUpdate') +
+    $.newMimirStatefulSet(name, 3, container, store_gateway_data_pvc) +
     statefulSet.mixin.spec.template.spec.withTerminationGracePeriodSeconds(120) +
-    // Parallelly scale up/down store-gateway instances instead of starting them
-    // one by one. This does NOT affect rolling updates: they will continue to be
-    // rolled out one by one (the next pod will be rolled out once the previous is
-    // ready).
-    statefulSet.mixin.spec.withPodManagementPolicy('Parallel') +
     $.util.configVolumeMount($._config.overrides_configmap, $._config.overrides_configmap_mountpoint) +
     (if with_anti_affinity then $.util.antiAffinity else {}),
 
@@ -91,11 +77,8 @@
     $.util.serviceFor($.store_gateway_statefulset, $._config.service_ignored_labels),
 
   store_gateway_pdb:
-    podDisruptionBudget.new() +
-    podDisruptionBudget.mixin.metadata.withName('store-gateway-pdb') +
-    podDisruptionBudget.mixin.metadata.withLabels({ name: 'store-gateway-pdb' }) +
-    podDisruptionBudget.mixin.spec.selector.withMatchLabels({ name: 'store-gateway' }) +
     // To avoid any disruption in the read path we need at least 1 replica of each
     // block available, so the disruption budget depends on the blocks replication factor.
-    podDisruptionBudget.mixin.spec.withMaxUnavailable(if $._config.store_gateway_replication_factor > 1 then $._config.store_gateway_replication_factor - 1 else 1),
+    local maxUnavailable = if $._config.store_gateway_replication_factor > 1 then $._config.store_gateway_replication_factor - 1 else 1;
+    $.newMimirPdb('store-gateway', maxUnavailable),
 }