Support for ReplicaSet targets (#3049)

* Added services and replicasets to mirrord ls output * Added replicaset variant * Updated schema and configuration.md * Changelog * Fix operator setup * Update mirrord/config/src/target.rs Co-authored-by: meowjesty <[email protected]> --------- Co-authored-by: meowjesty <[email protected]>
metalbear-co · Jan 30, 2025 · 2c4bf38 · 2c4bf38
1 parent 16d1dfe
commit 2c4bf38
Show file tree

Hide file tree

Showing 13 changed files with 252 additions and 15 deletions.
diff --git a/changelog.d/+replicaset-target.added.md b/changelog.d/+replicaset-target.added.md
@@ -0,0 +1 @@
+Added Kubernetes ReplicaSet as a new target type.
diff --git a/mirrord-schema.json b/mirrord-schema.json
@@ -1669,6 +1669,24 @@
         }
       ]
     },
+    "ReplicaSetTarget": {
+      "type": "object",
+      "required": [
+        "replica_set"
+      ],
+      "properties": {
+        "container": {
+          "type": [
+            "string",
+            "null"
+          ]
+        },
+        "replica_set": {
+          "type": "string"
+        }
+      },
+      "additionalProperties": false
+    },
     "RolloutTarget": {
       "description": "<!--${internal}--> Mirror the rollout specified by [`RolloutTarget::rollout`].",
       "type": "object",
@@ -1733,7 +1751,7 @@
       "additionalProperties": false
     },
     "Target": {
-      "description": "<!--${internal}--> ## path\n\nSpecifies the running pod (or deployment) to mirror.\n\nSupports: - `targetless` - `pod/{pod-name}[/container/{container-name}]`; - `deployment/{deployment-name}[/container/{container-name}]`; - `rollout/{rollout-name}[/container/{container-name}]`; - `job/{job-name}[/container/{container-name}]`; - `cronjob/{cronjob-name}[/container/{container-name}]`; - `statefulset/{statefulset-name}[/container/{container-name}]`; - `service/{service-name}[/container/{container-name}]`;",
+      "description": "<!--${internal}--> ## path\n\nSpecifies the running pod (or deployment) to mirror.\n\nSupports: - `targetless` - `pod/{pod-name}[/container/{container-name}]`; - `deployment/{deployment-name}[/container/{container-name}]`; - `rollout/{rollout-name}[/container/{container-name}]`; - `job/{job-name}[/container/{container-name}]`; - `cronjob/{cronjob-name}[/container/{container-name}]`; - `statefulset/{statefulset-name}[/container/{container-name}]`; - `service/{service-name}[/container/{container-name}]`; - `replicaset/{replicaset-name}[/container/{container-name}]`;",
       "anyOf": [
         {
           "description": "<!--${internal}--> [Deployment](https://kubernetes.io/docs/concepts/workloads/controllers/deployment/).",
@@ -1791,6 +1809,14 @@
             }
           ]
         },
+        {
+          "description": "<!--${internal}--> [ReplicaSet](https://kubernetes.io/docs/concepts/workloads/controllers/replicaset/).",
+          "allOf": [
+            {
+              "$ref": "#/definitions/ReplicaSetTarget"
+            }
+          ]
+        },
         {
           "description": "<!--${internal}--> Spawn a new pod.",
           "type": "null"

diff --git a/mirrord/cli/src/verify_config.rs b/mirrord/cli/src/verify_config.rs
@@ -8,8 +8,8 @@ use mirrord_config::{
     feature::FeatureConfig,
     target::{
         cron_job::CronJobTarget, deployment::DeploymentTarget, job::JobTarget, pod::PodTarget,
-        rollout::RolloutTarget, service::ServiceTarget, stateful_set::StatefulSetTarget, Target,
-        TargetConfig,
+        replica_set::ReplicaSetTarget, rollout::RolloutTarget, service::ServiceTarget,
+        stateful_set::StatefulSetTarget, Target, TargetConfig,
     },
 };
 use serde::Serialize;
@@ -47,6 +47,9 @@ enum VerifiedTarget {
 
     #[serde(untagged)]
     Service(ServiceTarget),
+
+    #[serde(untagged)]
+    ReplicaSet(ReplicaSetTarget),
 }
 
 impl From<Target> for VerifiedTarget {
@@ -59,6 +62,7 @@ impl From<Target> for VerifiedTarget {
             Target::CronJob(target) => Self::CronJob(target),
             Target::StatefulSet(target) => Self::StatefulSet(target),
             Target::Service(target) => Self::Service(target),
+            Target::ReplicaSet(target) => Self::ReplicaSet(target),
             Target::Targetless => Self::Targetless,
         }
     }
@@ -75,6 +79,7 @@ impl From<VerifiedTarget> for TargetType {
             VerifiedTarget::CronJob(_) => TargetType::CronJob,
             VerifiedTarget::StatefulSet(_) => TargetType::StatefulSet,
             VerifiedTarget::Service(_) => TargetType::Service,
+            VerifiedTarget::ReplicaSet(_) => TargetType::ReplicaSet,
         }
     }
 }
@@ -106,6 +111,7 @@ enum TargetType {
     CronJob,
     StatefulSet,
     Service,
+    ReplicaSet,
 }
 
 impl core::fmt::Display for TargetType {
@@ -119,6 +125,7 @@ impl core::fmt::Display for TargetType {
             TargetType::CronJob => "cronjob",
             TargetType::StatefulSet => "statefulset",
             TargetType::Service => "service",
+            TargetType::ReplicaSet => "replicaset",
         };
 
         f.write_str(stringifed)
@@ -136,6 +143,7 @@ impl TargetType {
             Self::CronJob,
             Self::StatefulSet,
             Self::Service,
+            Self::ReplicaSet,
         ]
         .into_iter()
     }
@@ -146,7 +154,7 @@ impl TargetType {
             Self::Pod => !(config.copy_target.enabled && config.copy_target.scale_down),
             Self::Job | Self::CronJob => config.copy_target.enabled,
             Self::Service => !config.copy_target.enabled,
-            Self::Deployment | Self::StatefulSet => true,
+            Self::Deployment | Self::StatefulSet | Self::ReplicaSet => true,
         }
     }
 }

diff --git a/mirrord/config/configuration.md b/mirrord/config/configuration.md
@@ -1652,6 +1652,7 @@ Supports:
 - `statefulset/{statefulset-name}[/container/{container-name}]`; (requires mirrord
   Operator)
 - `service/{service-name}[/container/{container-name}]`; (requires mirrord Operator)
+- `replicaset/{replicaset-name}[/container/{replicaset-name}]`; (requires mirrord Operator)
 
 ## telemetry {#root-telemetry}
 Controls whether or not mirrord sends telemetry data to MetalBear cloud.

diff --git a/mirrord/config/src/target.rs b/mirrord/config/src/target.rs
@@ -1,8 +1,8 @@
-use core::fmt;
-use std::str::FromStr;
+use std::{fmt, str::FromStr};
 
 use cron_job::CronJobTarget;
 use mirrord_analytics::CollectAnalytics;
+use replica_set::ReplicaSetTarget;
 use schemars::{gen::SchemaGenerator, schema::SchemaObject, JsonSchema};
 use serde::{Deserialize, Serialize};
 
@@ -23,6 +23,7 @@ pub mod cron_job;
 pub mod deployment;
 pub mod job;
 pub mod pod;
+pub mod replica_set;
 pub mod rollout;
 pub mod service;
 pub mod stateful_set;
@@ -147,6 +148,7 @@ pub struct TargetConfig {
     /// - `statefulset/{statefulset-name}[/container/{container-name}]`; (requires mirrord
     ///   Operator)
     /// - `service/{service-name}[/container/{container-name}]`; (requires mirrord Operator)
+    /// - `replicaset/{replicaset-name}[/container/{container-name}]`; (requires mirrord Operator)
     #[serde(skip_serializing_if = "Option::is_none")]
     pub path: Option<Target>,
 
@@ -221,6 +223,7 @@ mirrord-layer failed to parse the provided target!
     >> `cronjob/{cronjob-name}[/container/{container-name}]`;
     >> `statefulset/{statefulset-name}[/container/{container-name}]`;
     >> `service/{service-name}[/container/{container-name}]`;
+    >> `replicaset/{replicaset-name}[/container/{container-name}]`;
 
 - Note:
     >> specifying container name is optional, defaults to a container chosen by mirrord
@@ -246,6 +249,7 @@ mirrord-layer failed to parse the provided target!
 /// - `cronjob/{cronjob-name}[/container/{container-name}]`;
 /// - `statefulset/{statefulset-name}[/container/{container-name}]`;
 /// - `service/{service-name}[/container/{container-name}]`;
+/// - `replicaset/{replicaset-name}[/container/{container-name}]`;
 #[warn(clippy::wildcard_enum_match_arm)]
 #[derive(Serialize, Deserialize, Clone, Eq, PartialEq, Hash, Debug, JsonSchema)]
 #[serde(untagged, deny_unknown_fields)]
@@ -282,6 +286,10 @@ pub enum Target {
     /// [Service](https://kubernetes.io/docs/concepts/services-networking/service/).
     Service(service::ServiceTarget),
 
+    /// <!--${internal}-->
+    /// [ReplicaSet](https://kubernetes.io/docs/concepts/workloads/controllers/replicaset/).
+    ReplicaSet(replica_set::ReplicaSetTarget),
+
     /// <!--${internal}-->
     /// Spawn a new pod.
     Targetless,
@@ -305,6 +313,7 @@ impl FromStr for Target {
             Some("cronjob") => cron_job::CronJobTarget::from_split(&mut split).map(Target::CronJob),
             Some("statefulset") => stateful_set::StatefulSetTarget::from_split(&mut split).map(Target::StatefulSet),
             Some("service") => service::ServiceTarget::from_split(&mut split).map(Target::Service),
+            Some("replicaset") => replica_set::ReplicaSetTarget::from_split(&mut split).map(Target::ReplicaSet),
             _ => Err(ConfigError::InvalidTarget(format!(
                 "Provided target: {target} is unsupported. Did you remember to add a prefix, e.g. pod/{target}? \n{FAIL_PARSE_DEPLOYMENT_OR_POD}",
             ))),
@@ -323,6 +332,7 @@ impl Target {
             Target::CronJob(target) => target.cron_job.clone(),
             Target::StatefulSet(target) => target.stateful_set.clone(),
             Target::Service(target) => target.service.clone(),
+            Target::ReplicaSet(target) => target.replica_set.clone(),
             Target::Targetless => {
                 unreachable!("this shouldn't happen - called from operator on a flow where it's not targetless.")
             }
@@ -338,7 +348,11 @@ impl Target {
     pub(super) fn requires_operator(&self) -> bool {
         matches!(
             self,
-            Target::Job(_) | Target::CronJob(_) | Target::StatefulSet(_) | Target::Service(_)
+            Target::Job(_)
+                | Target::CronJob(_)
+                | Target::StatefulSet(_)
+                | Target::Service(_)
+                | Target::ReplicaSet(_)
         )
     }
 }
@@ -399,6 +413,7 @@ impl_target_display!(JobTarget, job, "job");
 impl_target_display!(CronJobTarget, cron_job, "cronjob");
 impl_target_display!(StatefulSetTarget, stateful_set, "statefulset");
 impl_target_display!(ServiceTarget, service, "service");
+impl_target_display!(ReplicaSetTarget, replica_set, "replicaset");
 
 impl fmt::Display for Target {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
@@ -411,6 +426,7 @@ impl fmt::Display for Target {
             Target::CronJob(target) => target.fmt(f),
             Target::StatefulSet(target) => target.fmt(f),
             Target::Service(target) => target.fmt(f),
+            Target::ReplicaSet(target) => target.fmt(f),
         }
     }
 }
@@ -426,6 +442,7 @@ impl TargetDisplay for Target {
             Target::CronJob(target) => target.type_(),
             Target::StatefulSet(target) => target.type_(),
             Target::Service(target) => target.type_(),
+            Target::ReplicaSet(target) => target.type_(),
         }
     }
 
@@ -439,6 +456,7 @@ impl TargetDisplay for Target {
             Target::CronJob(target) => target.name(),
             Target::StatefulSet(target) => target.name(),
             Target::Service(target) => target.name(),
+            Target::ReplicaSet(target) => target.name(),
         }
     }
 
@@ -452,6 +470,7 @@ impl TargetDisplay for Target {
             Target::CronJob(target) => target.container(),
             Target::StatefulSet(target) => target.container(),
             Target::Service(target) => target.container(),
+            Target::ReplicaSet(target) => target.container(),
         }
     }
 }
@@ -469,6 +488,7 @@ bitflags::bitflags! {
         const CRON_JOB = 64;
         const STATEFUL_SET = 128;
         const SERVICE = 256;
+        const REPLICA_SET = 512;
     }
 }
 
@@ -522,6 +542,12 @@ impl CollectAnalytics for &TargetConfig {
                         flags |= TargetAnalyticFlags::CONTAINER;
                     }
                 }
+                Target::ReplicaSet(target) => {
+                    flags |= TargetAnalyticFlags::REPLICA_SET;
+                    if target.container.is_some() {
+                        flags |= TargetAnalyticFlags::CONTAINER;
+                    }
+                }
                 Target::Targetless => {
                     // Targetless is essentially 0, so no need to set any flags.
                 }

diff --git a/mirrord/config/src/target/replica_set.rs b/mirrord/config/src/target/replica_set.rs
@@ -0,0 +1,35 @@
+use std::str::Split;
+
+use schemars::JsonSchema;
+use serde::{Deserialize, Serialize};
+
+use super::{FromSplit, FAIL_PARSE_DEPLOYMENT_OR_POD};
+use crate::config::{ConfigError, Result};
+
+#[derive(Serialize, Deserialize, Clone, Eq, PartialEq, Hash, Debug, JsonSchema)]
+#[serde(deny_unknown_fields)]
+pub struct ReplicaSetTarget {
+    pub replica_set: String,
+    pub container: Option<String>,
+}
+
+impl FromSplit for ReplicaSetTarget {
+    fn from_split(split: &mut Split<char>) -> Result<Self> {
+        let replica_set = split
+            .next()
+            .ok_or_else(|| ConfigError::InvalidTarget(FAIL_PARSE_DEPLOYMENT_OR_POD.to_string()))?;
+        match (split.next(), split.next()) {
+            (Some("container"), Some(container)) => Ok(Self {
+                replica_set: replica_set.to_string(),
+                container: Some(container.to_string()),
+            }),
+            (None, None) => Ok(Self {
+                replica_set: replica_set.to_string(),
+                container: None,
+            }),
+            _ => Err(ConfigError::InvalidTarget(
+                FAIL_PARSE_DEPLOYMENT_OR_POD.to_string(),
+            )),
+        }
+    }
+}
diff --git a/mirrord/kube/src/api/kubernetes/seeker.rs b/mirrord/kube/src/api/kubernetes/seeker.rs
@@ -3,9 +3,9 @@ use std::fmt;
 use futures::{stream, Stream, StreamExt, TryStreamExt};
 use k8s_openapi::{
     api::{
-        apps::v1::{Deployment, StatefulSet},
+        apps::v1::{Deployment, ReplicaSet, StatefulSet},
         batch::v1::{CronJob, Job},
-        core::v1::Pod,
+        core::v1::{Pod, Service},
     },
     ClusterResourceScope, Metadata, NamespaceResourceScope,
 };
@@ -42,16 +42,25 @@ impl KubeResourceSeeker<'_> {
             .collect())
     }
 
-    /// Returns all resource types ie. [`Pod`], [`Deployment`], [`Rollout`], [`Job`], [`CronJob`],
-    /// and [`StatefulSet`]
+    /// Returns all resource of all types:
+    /// 1. [`Deployment`]s,
+    /// 2. [`Rollout`]s,
+    /// 3. [`StatefulSet`]s
+    /// 4. [`CronJob`]s
+    /// 5. [`Job`]s
+    /// 6. [`Service`]s
+    /// 7. [`ReplicaSet`]s
+    /// 8. [`Pod`]s
     pub async fn all(&self) -> Result<Vec<String>> {
-        let (pods, deployments, rollouts, jobs, cronjobs, statefulsets) = tokio::try_join!(
+        let (pods, deployments, rollouts, jobs, cronjobs, statefulsets, services, replicasets) = tokio::try_join!(
             self.pods(),
             self.simple_list_resource::<Deployment>("deployment"),
             self.simple_list_resource::<Rollout>("rollout"),
             self.simple_list_resource::<Job>("job"),
             self.simple_list_resource::<CronJob>("cronjob"),
             self.simple_list_resource::<StatefulSet>("statefulset"),
+            self.simple_list_resource::<Service>("service"),
+            self.simple_list_resource::<ReplicaSet>("replicaset"),
         )?;
 
         Ok(deployments
@@ -60,6 +69,8 @@ impl KubeResourceSeeker<'_> {
             .chain(statefulsets)
             .chain(cronjobs)
             .chain(jobs)
+            .chain(services)
+            .chain(replicasets)
             .chain(pods)
             .collect())
     }

diff --git a/mirrord/kube/src/api/runtime.rs b/mirrord/kube/src/api/runtime.rs
@@ -33,6 +33,7 @@ pub mod cron_job;
 pub mod deployment;
 pub mod job;
 pub mod pod;
+pub mod replica_set;
 pub mod rollout;
 pub mod service;
 pub mod stateful_set;
@@ -374,6 +375,7 @@ impl RuntimeDataProvider for Target {
             Target::CronJob(target) => target.runtime_data(client, namespace).await,
             Target::StatefulSet(target) => target.runtime_data(client, namespace).await,
             Target::Service(target) => target.runtime_data(client, namespace).await,
+            Target::ReplicaSet(target) => target.runtime_data(client, namespace).await,
             Target::Targetless => Err(KubeApiError::MissingRuntimeData),
         }
     }
@@ -389,6 +391,7 @@ impl RuntimeDataProvider for ResolvedTarget<true> {
             Self::CronJob(target) => target.runtime_data(client, namespace).await,
             Self::StatefulSet(target) => target.runtime_data(client, namespace).await,
             Self::Service(target) => target.runtime_data(client, namespace).await,
+            Self::ReplicaSet(target) => target.runtime_data(client, namespace).await,
             Self::Targetless(_) => Err(KubeApiError::MissingRuntimeData),
         }
     }
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1 @@
		Added Kubernetes ReplicaSet as a new target type.