ref(k8s): Improve update behavior of pods in the replicator statefulset (supabase#284)

Author: iambriccardo

etl-api/Cargo.toml

@@ -23,6 +23,7 @@ async-trait = { workspace = true }
 aws-lc-rs = { workspace = true, features = ["alloc", "aws-lc-sys"] }
 base64 = { workspace = true, features = ["std"] }
 constant_time_eq = { workspace = true }
+chrono = { workspace = true }
 k8s-openapi = { workspace = true, features = ["latest"] }
 kube = { workspace = true, features = [
     "runtime",

etl-api/src/k8s_client.rs

@@ -5,6 +5,7 @@ use k8s_openapi::api::{
     core::v1::{ConfigMap, Pod, Secret},
 };
 use serde_json::json;
+use std::collections::BTreeMap;
 use thiserror::Error;
 use tracing::info;
 
@@ -75,6 +76,7 @@ pub trait K8sClient: Send + Sync {
         &self,
         prefix: &str,
         replicator_image: &str,
+        template_annotations: Option<BTreeMap<String, String>>,
     ) -> Result<(), K8sError>;
 
     async fn delete_stateful_set(&self, prefix: &str) -> Result<(), K8sError>;
@@ -113,6 +115,7 @@ pub const TRUSTED_ROOT_CERT_CONFIG_MAP_NAME: &str = "trusted-root-certs-config";
 pub const TRUSTED_ROOT_CERT_KEY_NAME: &str = "trusted_root_certs";
 const PG_PASSWORD_ENV_VAR_NAME: &str = "APP_PIPELINE__PG_CONNECTION__PASSWORD";
 const BIG_QUERY_SA_KEY_ENV_VAR_NAME: &str = "APP_DESTINATION__BIG_QUERY__SERVICE_ACCOUNT_KEY";
+pub const RESTARTED_AT_ANNOTATION_KEY: &str = "etl.supabase.com/restarted-at";
 
 impl HttpK8sClient {
     pub async fn new() -> Result<HttpK8sClient, K8sError> {
@@ -304,6 +307,7 @@ impl K8sClient for HttpK8sClient {
         &self,
         prefix: &str,
         replicator_image: &str,
+        template_annotations: Option<BTreeMap<String, String>>,
     ) -> Result<(), K8sError> {
         info!("patching stateful set");
 
@@ -315,7 +319,7 @@ 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 stateful_set_json = json!({
+        let mut stateful_set_json = json!({
           "apiVersion": "apps/v1",
           "kind": "StatefulSet",
           "metadata": {
@@ -447,15 +451,27 @@ impl K8sClient for HttpK8sClient {
           }
         });
 
+        // Attach template annotations (e.g., restart checksum) to trigger a rolling restart
+        if let Some(annotations) = template_annotations
+            && let Some(template) = stateful_set_json
+                .get_mut("spec")
+                .and_then(|s| s.get_mut("template"))
+                .and_then(|t| t.get_mut("metadata"))
+        {
+            // Insert annotations map
+            let annotations_value = serde_json::to_value(annotations)?;
+            if let Some(obj) = template.as_object_mut() {
+                obj.insert("annotations".to_string(), annotations_value);
+            }
+        }
+
         let stateful_set: StatefulSet = serde_json::from_value(stateful_set_json)?;
 
         let pp = PatchParams::apply(&stateful_set_name);
         self.stateful_sets_api
             .patch(&stateful_set_name, &pp, &Patch::Apply(stateful_set))
             .await?;
 
-        self.delete_pod(prefix).await?;
-
         info!("patched stateful set");
 
         Ok(())
@@ -477,7 +493,7 @@ impl K8sClient for HttpK8sClient {
                 e => return Err(e.into()),
             },
         }
-        self.delete_pod(prefix).await?;
+
         info!("deleted stateful set");
 
         Ok(())

etl-api/src/routes/pipelines.rs

@@ -4,12 +4,14 @@ use actix_web::{
     post,
     web::{Data, Json, Path},
 };
+use chrono::Utc;
 use etl_config::shared::{ReplicatorConfig, SupabaseConfig, TlsConfig};
 use etl_postgres::replication::{TableLookupError, get_table_name_from_oid, health, state};
 use etl_postgres::schema::TableId;
 use secrecy::ExposeSecret;
 use serde::{Deserialize, Serialize};
 use sqlx::{PgPool, PgTransaction};
+use std::collections::BTreeMap;
 use std::ops::DerefMut;
 use thiserror::Error;
 use utoipa::ToSchema;
@@ -24,8 +26,8 @@ use crate::db::images::{Image, ImagesDbError};
 use crate::db::pipelines::{Pipeline, PipelinesDbError};
 use crate::db::replicators::{Replicator, ReplicatorsDbError};
 use crate::db::sources::{Source, SourcesDbError, source_exists};
-use crate::k8s_client::TRUSTED_ROOT_CERT_KEY_NAME;
 use crate::k8s_client::{K8sClient, K8sError, PodPhase, TRUSTED_ROOT_CERT_CONFIG_MAP_NAME};
+use crate::k8s_client::{RESTARTED_AT_ANNOTATION_KEY, TRUSTED_ROOT_CERT_KEY_NAME};
 use crate::routes::{
     ErrorMessage, TenantIdError, connect_to_source_database_with_defaults, extract_tenant_id,
 };
@@ -360,7 +362,6 @@ pub enum PipelineStatus {
     Stopped,
     Starting,
     Started,
-    Stopping,
     Unknown,
     Failed,
 }
@@ -1088,7 +1089,7 @@ pub async fn update_pipeline_config(
     Ok(Json(response))
 }
 
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
 struct Secrets {
     postgres_password: String,
     big_query_service_account_key: Option<String>,
@@ -1107,7 +1108,7 @@ async fn create_or_update_pipeline_in_k8s(
 
     // We create the secrets.
     let secrets = build_secrets(&source.config, &destination.config);
-    create_or_update_secrets(k8s_client, &prefix, secrets).await?;
+    create_or_update_secrets(k8s_client, &prefix, secrets.clone()).await?;
 
     // We create the replicator configuration.
     let replicator_config = build_replicator_config(
@@ -1120,14 +1121,31 @@ async fn create_or_update_pipeline_in_k8s(
         },
     )
     .await?;
-    create_or_update_config(k8s_client, &prefix, replicator_config).await?;
+    create_or_update_config(k8s_client, &prefix, replicator_config.clone()).await?;
+
+    // To force restart everytime, we want to annotate the stateful set with the current UTC time for every
+    // start call. Technically we can optimally perform a restart by calculating a checksum on a deterministic
+    // set of inputs like the configs, states in the database, etc... however we deemed that too cumbersome
+    // and risky, since forgetting a component will lead to the pipeline not restarting.
+    let mut annotations = BTreeMap::new();
+    annotations.insert(
+        RESTARTED_AT_ANNOTATION_KEY.to_string(),
+        get_restarted_at_annotation_value(),
+    );
 
     // We create the replicator stateful set.
-    create_or_update_replicator(k8s_client, &prefix, image.name).await?;
+    create_or_update_replicator(k8s_client, &prefix, image.name, Some(annotations)).await?;
 
     Ok(())
 }
 
+fn get_restarted_at_annotation_value() -> String {
+    let now = Utc::now();
+    // We use nanoseconds to decrease the likelihood of generating the same annotation in sequence,
+    // which would not result in a restart.
+    now.to_rfc3339_opts(chrono::SecondsFormat::Nanos, true)
+}
+
 async fn delete_pipeline_in_k8s(
     k8s_client: &dyn K8sClient,
     tenant_id: &str,
@@ -1281,9 +1299,10 @@ async fn create_or_update_replicator(
     k8s_client: &dyn K8sClient,
     prefix: &str,
     replicator_image: String,
+    template_annotations: Option<BTreeMap<String, String>>,
 ) -> Result<(), PipelineError> {
     k8s_client
-        .create_or_update_stateful_set(prefix, &replicator_image)
+        .create_or_update_stateful_set(prefix, &replicator_image, template_annotations)
         .await?;
 
     Ok(())

etl-api/tests/common/k8s_client.rs

@@ -67,6 +67,7 @@ impl K8sClient for MockK8sClient {
         &self,
         _prefix: &str,
         _replicator_image: &str,
+        _template_annotations: Option<BTreeMap<String, String>>,
     ) -> Result<(), K8sError> {
         Ok(())
     }