feat(k8s): Improve resource tuning of K8S (#345)

Author: iambriccardo

etl-api/src/k8s/base.rs

@@ -0,0 +1,126 @@
+use async_trait::async_trait;
+use k8s_openapi::api::core::v1::ConfigMap;
+use std::collections::BTreeMap;
+use thiserror::Error;
+
+/// Errors emitted by the Kubernetes integration.
+///
+/// Variants wrap lower-level libraries where appropriate to preserve context.
+#[derive(Debug, Error)]
+pub enum K8sError {
+    /// A serialization or deserialization error while building or parsing
+    /// Kubernetes resources.
+    #[error("An error occurred in serde: {0}")]
+    Serde(#[from] serde_json::error::Error),
+    /// An error returned by the [`kube`] client when talking to the API
+    /// server.
+    #[error("An error occurred with kube: {0}")]
+    Kube(#[from] kube::Error),
+    /// The environment-dependent replicator configuration could not be
+    /// determined.
+    #[error("An error occurred while configuring the replicator")]
+    ReplicatorConfiguration,
+}
+
+/// A simplified view of a pod phase.
+///
+/// This mirrors the string phases reported by Kubernetes but only tracks the
+/// states needed by the API. Unknown values map to [`PodPhase::Unknown`].
+pub enum PodPhase {
+    Pending,
+    Running,
+    Succeeded,
+    Failed,
+    Unknown,
+}
+
+impl From<&str> for PodPhase {
+    /// Converts a Kubernetes pod phase string into a [`PodPhase`].
+    ///
+    /// Unrecognized values result in [`PodPhase::Unknown`].
+    fn from(value: &str) -> Self {
+        match value {
+            "Pending" => PodPhase::Pending,
+            "Running" => PodPhase::Running,
+            "Succeeded" => PodPhase::Succeeded,
+            "Failed" => PodPhase::Failed,
+            _ => PodPhase::Unknown,
+        }
+    }
+}
+
+/// Client interface describing the Kubernetes operations used by the API.
+///
+/// Implementations are expected to be idempotent where possible by issuing
+/// server-side apply patches for create-or-update behaviors.
+#[async_trait]
+pub trait K8sClient: Send + Sync {
+    /// Creates or updates the Postgres password secret for a replicator.
+    ///
+    /// The secret name is derived from `prefix` and is stored in the
+    /// data-plane namespace.
+    async fn create_or_update_postgres_secret(
+        &self,
+        prefix: &str,
+        postgres_password: &str,
+    ) -> Result<(), K8sError>;
+
+    /// Creates or updates the BigQuery service account secret for a
+    /// replicator.
+    ///
+    /// The secret name is derived from `prefix` and is stored in the
+    /// data-plane namespace.
+    async fn create_or_update_bq_secret(
+        &self,
+        prefix: &str,
+        bq_service_account_key: &str,
+    ) -> Result<(), K8sError>;
+
+    /// Deletes the Postgres password secret for a replicator if it exists.
+    async fn delete_postgres_secret(&self, prefix: &str) -> Result<(), K8sError>;
+
+    /// Deletes the BigQuery service account secret for a replicator if it
+    /// exists.
+    async fn delete_bq_secret(&self, prefix: &str) -> Result<(), K8sError>;
+
+    /// Retrieves a named [`ConfigMap`].
+    async fn get_config_map(&self, config_map_name: &str) -> Result<ConfigMap, K8sError>;
+
+    /// Creates or updates the replicator configuration [`ConfigMap`].
+    ///
+    /// The config map stores two YAML documents: a base and a production
+    /// override.
+    async fn create_or_update_config_map(
+        &self,
+        prefix: &str,
+        base_config: &str,
+        prod_config: &str,
+    ) -> Result<(), K8sError>;
+
+    /// Deletes the replicator configuration [`ConfigMap`] if it exists.
+    async fn delete_config_map(&self, prefix: &str) -> Result<(), K8sError>;
+
+    /// Creates or updates the replicator [`StatefulSet`].
+    ///
+    /// The set references previously created secrets and config maps. Optional
+    /// `template_annotations` may be used to trigger a rolling restart.
+    async fn create_or_update_stateful_set(
+        &self,
+        prefix: &str,
+        replicator_image: &str,
+        template_annotations: Option<BTreeMap<String, String>>,
+    ) -> Result<(), K8sError>;
+
+    /// Deletes the replicator [`StatefulSet`] if it exists.
+    async fn delete_stateful_set(&self, prefix: &str) -> Result<(), K8sError>;
+
+    /// Returns the phase of the replicator pod.
+    async fn get_pod_phase(&self, prefix: &str) -> Result<PodPhase, K8sError>;
+
+    /// Reports whether the replicator container terminated with a non-zero exit
+    /// code.
+    async fn has_replicator_container_error(&self, prefix: &str) -> Result<bool, K8sError>;
+
+    /// Deletes the replicator pod if it exists.
+    async fn delete_pod(&self, prefix: &str) -> Result<(), K8sError>;
+}

etl-api/src/k8s_client.rs renamed to etl-api/src/k8s/http.rs

@@ -1,124 +1,114 @@
+use crate::k8s::{K8sClient, K8sError, PodPhase};
 use async_trait::async_trait;
 use base64::{Engine, prelude::BASE64_STANDARD};
+use etl_config::Environment;
 use k8s_openapi::api::{
     apps::v1::StatefulSet,
     core::v1::{ConfigMap, Pod, Secret},
 };
-use serde_json::json;
-use std::collections::BTreeMap;
-use thiserror::Error;
-use tracing::info;
-
 use kube::{
     Client,
     api::{Api, DeleteParams, Patch, PatchParams},
 };
+use serde_json::json;
+use std::collections::BTreeMap;
+use tracing::info;
 
-#[derive(Debug, Error)]
-pub enum K8sError {
-    #[error("serde_json error: {0}")]
-    Serde(#[from] serde_json::error::Error),
-
-    #[error("kube error: {0}")]
-    Kube(#[from] kube::Error),
-}
-
-pub enum PodPhase {
-    Pending,
-    Running,
-    Succeeded,
-    Failed,
-    Unknown,
-}
-
-impl From<&str> for PodPhase {
-    fn from(value: &str) -> Self {
-        match value {
-            "Pending" => PodPhase::Pending,
-            "Running" => PodPhase::Running,
-            "Succeeded" => PodPhase::Succeeded,
-            "Failed" => PodPhase::Failed,
-            _ => PodPhase::Unknown,
-        }
-    }
-}
-
-#[async_trait]
-pub trait K8sClient: Send + Sync {
-    async fn create_or_update_postgres_secret(
-        &self,
-        prefix: &str,
-        postgres_password: &str,
-    ) -> Result<(), K8sError>;
-
-    async fn create_or_update_bq_secret(
-        &self,
-        prefix: &str,
-        bq_service_account_key: &str,
-    ) -> Result<(), K8sError>;
-
-    async fn delete_postgres_secret(&self, prefix: &str) -> Result<(), K8sError>;
-
-    async fn delete_bq_secret(&self, prefix: &str) -> Result<(), K8sError>;
-
-    async fn get_config_map(&self, config_map_name: &str) -> Result<ConfigMap, K8sError>;
-
-    async fn create_or_update_config_map(
-        &self,
-        prefix: &str,
-        base_config: &str,
-        prod_config: &str,
-    ) -> Result<(), K8sError>;
-
-    async fn delete_config_map(&self, prefix: &str) -> Result<(), K8sError>;
-
-    async fn create_or_update_stateful_set(
-        &self,
-        prefix: &str,
-        replicator_image: &str,
-        template_annotations: Option<BTreeMap<String, String>>,
-    ) -> Result<(), K8sError>;
-
-    async fn delete_stateful_set(&self, prefix: &str) -> Result<(), K8sError>;
-
-    async fn get_pod_phase(&self, prefix: &str) -> Result<PodPhase, K8sError>;
-
-    async fn has_replicator_container_error(&self, prefix: &str) -> Result<bool, K8sError>;
-
-    async fn delete_pod(&self, prefix: &str) -> Result<(), K8sError>;
-}
-
-#[derive(Debug)]
-pub struct HttpK8sClient {
-    secrets_api: Api<Secret>,
-    config_maps_api: Api<ConfigMap>,
-    stateful_sets_api: Api<StatefulSet>,
-    pods_api: Api<Pod>,
-}
-
+/// Secret name suffix for the BigQuery service account key.
 const BQ_SECRET_NAME_SUFFIX: &str = "bq-service-account-key";
+/// Secret name suffix for the Postgres password.
 const POSTGRES_SECRET_NAME_SUFFIX: &str = "postgres-password";
+/// ConfigMap name suffix for the replicator configuration files.
 const REPLICATOR_CONFIG_MAP_NAME_SUFFIX: &str = "replicator-config";
+/// StatefulSet name suffix for the replicator workload.
 const REPLICATOR_STATEFUL_SET_SUFFIX: &str = "replicator-stateful-set";
+/// Application label suffix used to group resources.
 const REPLICATOR_APP_SUFFIX: &str = "replicator-app";
+/// Container name suffix for the replicator container.
 const REPLICATOR_CONTAINER_NAME_SUFFIX: &str = "replicator";
+/// Container name suffix for the Vector sidecar.
 const VECTOR_CONTAINER_NAME_SUFFIX: &str = "vector";
+/// Namespace where data-plane resources are created.
 const DATA_PLANE_NAMESPACE: &str = "etl-data-plane";
+/// Secret storing the Logflare API key.
 const LOGFLARE_SECRET_NAME: &str = "replicator-logflare-api-key";
+/// Docker image used for the Vector sidecar.
 const VECTOR_IMAGE_NAME: &str = "timberio/vector:0.46.1-distroless-libc";
+/// ConfigMap name containing the Vector configuration.
 const VECTOR_CONFIG_MAP_NAME: &str = "replicator-vector-config";
+/// Volume name for the replicator config file.
 const REPLICATOR_CONFIG_FILE_VOLUME_NAME: &str = "replicator-config-file";
+/// Volume name for the Vector config file.
 const VECTOR_CONFIG_FILE_VOLUME_NAME: &str = "vector-config-file";
+/// Secret storing the Sentry DSN.
 const SENTRY_DSN_SECRET_NAME: &str = "replicator-sentry-dsn";
+/// EmptyDir volume name used to share logs.
 const LOGS_VOLUME_NAME: &str = "logs";
+/// ConfigMap name providing trusted root certificates.
 pub const TRUSTED_ROOT_CERT_CONFIG_MAP_NAME: &str = "trusted-root-certs-config";
+/// Key inside the trusted root certificates ConfigMap.
 pub const TRUSTED_ROOT_CERT_KEY_NAME: &str = "trusted_root_certs";
+/// Environment variable for the Postgres password.
 const PG_PASSWORD_ENV_VAR_NAME: &str = "APP_PIPELINE__PG_CONNECTION__PASSWORD";
+/// Environment variable for the BigQuery service account key.
 const BIG_QUERY_SA_KEY_ENV_VAR_NAME: &str = "APP_DESTINATION__BIG_QUERY__SERVICE_ACCOUNT_KEY";
+/// Pod template annotation used to trigger rolling restarts.
 pub const RESTARTED_AT_ANNOTATION_KEY: &str = "etl.supabase.com/restarted-at";
+/// Label used to identify replicator pods.
 const REPLICATOR_APP_LABEL: &str = "etl-replicator-app";
 
+/// Replicator memory limit tuned for `c6in.4xlarge` instances.
+const REPLICATOR_MAX_MEMORY_PROD: &str = "500Mi";
+/// Replicator CPU limit tuned for `c6in.4xlarge` instances.
+const REPLICATOR_MAX_CPU_PROD: &str = "100m";
+/// Replicator memory limit tuned for `t3.small` instances.
+const REPLICATOR_MAX_MEMORY_STAGING: &str = "100Mi";
+/// Replicator CPU limit tuned for `t3.small` instances.
+const REPLICATOR_MAX_CPU_STAGING: &str = "100m";
+
+/// Runtime limits derived from the current environment.
+struct DynamicReplicatorConfig {
+    max_memory: &'static str,
+    max_cpu: &'static str,
+}
+
+impl DynamicReplicatorConfig {
+    /// Loads the runtime limits for the current environment.
+    fn load() -> Result<Self, K8sError> {
+        let environment = Environment::load().map_err(|_| K8sError::ReplicatorConfiguration)?;
+
+        let config = match environment {
+            Environment::Prod => Self {
+                max_memory: REPLICATOR_MAX_MEMORY_PROD,
+                max_cpu: REPLICATOR_MAX_CPU_PROD,
+            },
+            _ => Self {
+                max_memory: REPLICATOR_MAX_MEMORY_STAGING,
+                max_cpu: REPLICATOR_MAX_CPU_STAGING,
+            },
+        };
+
+        Ok(config)
+    }
+}
+
+/// HTTP-based implementation of [`K8sClient`].
+///
+/// The client is namespaced to the data-plane namespace and uses server-side
+/// apply to keep resources in sync.
+#[derive(Debug)]
+pub struct HttpK8sClient {
+    secrets_api: Api<Secret>,
+    config_maps_api: Api<ConfigMap>,
+    stateful_sets_api: Api<StatefulSet>,
+    pods_api: Api<Pod>,
+}
+
 impl HttpK8sClient {
+    /// Creates a new [`HttpK8sClient`] using the ambient Kubernetes config.
+    ///
+    /// Prefers in-cluster configuration and falls back to the local kubeconfig
+    /// when running outside the cluster.
     pub async fn new() -> Result<HttpK8sClient, K8sError> {
         let client = Client::try_default().await?;
 
@@ -320,6 +310,8 @@ impl K8sClient for HttpK8sClient {
         let bq_secret_name = format!("{prefix}-{BQ_SECRET_NAME_SUFFIX}");
         let replicator_config_map_name = format!("{prefix}-{REPLICATOR_CONFIG_MAP_NAME_SUFFIX}");
 
+        let config = DynamicReplicatorConfig::load()?;
+
         let mut stateful_set_json = json!({
           "apiVersion": "apps/v1",
           "kind": "StatefulSet",
@@ -393,11 +385,12 @@ impl K8sClient for HttpK8sClient {
                     ],
                     "resources": {
                       "limits": {
-                        "memory": "200Mi",
+                        "memory": config.max_memory,
+                        "cpu": config.max_cpu,
                       },
                       "requests": {
-                        "memory": "200Mi",
-                        "cpu": "100m"
+                        "memory": config.max_memory,
+                        "cpu": config.max_cpu,
                       }
                     },
                     "volumeMounts": [

etl-api/src/k8s/mod.rs

@@ -0,0 +1,18 @@
+//! Kubernetes integration for the ETL API.
+//!
+//! This module contains the abstractions and implementations used by the HTTP
+//! API to manage Kubernetes resources required by replicators (secrets, config
+//! maps, stateful sets, and pods). Consumers should depend on the trait
+//! [`K8sClient`] and avoid relying on a specific transport.
+//!
+//! The default client, [`http::HttpK8sClient`], is backed by the [`kube`]
+//! crate and talks to the cluster using the ambient configuration (in-cluster
+//! or local `~/.kube/config`). Keeping the abstraction in [`base`] lets us
+//! swap implementations in tests and non-Kubernetes environments.
+//!
+//! See [`base`] for errors, pod phase mapping, and the client trait.
+
+mod base;
+pub mod http;
+
+pub use base::*;

etl-api/src/lib.rs

@@ -1,14 +1,8 @@
-//! ETL API service for managing data replication pipelines.
-//!
-//! Provides a REST API for configuring and managing ETL pipelines, including tenants,
-//! sources, destinations, and replication monitoring. Includes authentication, encryption,
-//! Kubernetes integration, and comprehensive OpenAPI documentation.
-
 pub mod authentication;
 pub mod config;
 pub mod configs;
 pub mod db;
-pub mod k8s_client;
+pub mod k8s;
 pub mod routes;
 pub mod span_builder;
 pub mod startup;