From 46538b7eda32fe0dfcab47803a2a81761d204c97 Mon Sep 17 00:00:00 2001
From: Arthur Depasse <arthur.depasse@qube-rt.com>
Date: Tue, 10 Mar 2026 10:53:52 +0000
Subject: [PATCH] Allow configuring the PIA canary pods image in chart values

---
 charts/restate-operator-helm/templates/deployment.yaml | 2 ++
 charts/restate-operator-helm/values.yaml               | 4 ++++
 src/controllers/mod.rs                                 | 5 +++++
 src/controllers/restatecluster/controller.rs           | 3 +++
 src/controllers/restatecluster/reconcilers/compute.rs  | 4 +++-
 src/main.rs                                            | 9 +++++++++
 6 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/charts/restate-operator-helm/templates/deployment.yaml b/charts/restate-operator-helm/templates/deployment.yaml
index 0fdc653..8e48c72 100644
--- a/charts/restate-operator-helm/templates/deployment.yaml
+++ b/charts/restate-operator-helm/templates/deployment.yaml
@@ -60,6 +60,8 @@ spec:
             - name: AWS_POD_IDENTITY_ASSOCIATION_CLUSTER
               value: {{ .Values.awsPodIdentityAssociationCluster }}
           {{- end }}
+            - name: PIA_CANARY_IMAGE
+              value: {{ .Values.piaCanaryImage.repository }}:{{ .Values.piaCanaryImage.tag }}
           {{- with .Values.env }}
             {{- toYaml . | nindent 12 }}
           {{- end }}
diff --git a/charts/restate-operator-helm/values.yaml b/charts/restate-operator-helm/values.yaml
index d72bfe8..05d8851 100644
--- a/charts/restate-operator-helm/values.yaml
+++ b/charts/restate-operator-helm/values.yaml
@@ -15,6 +15,10 @@ podAnnotations: {}
 
 awsPodIdentityAssociationCluster: null
 
+piaCanaryImage:
+  repository: busybox
+  tag: uclibc
+
 podSecurityContext:
   fsGroup: 2000
   fsGroupChangePolicy: "OnRootMismatch"
diff --git a/src/controllers/mod.rs b/src/controllers/mod.rs
index 3a1a330..0ed5822 100644
--- a/src/controllers/mod.rs
+++ b/src/controllers/mod.rs
@@ -48,6 +48,9 @@ pub struct State {
 
     /// The default image to use for tunnel client pods
     tunnel_client_default_image: String,
+
+    /// The image to use for PIA canary pods
+    pub pia_canary_image: String,
 }
 
 /// State wrapper around the controller outputs for the web server
@@ -58,6 +61,7 @@ impl State {
         operator_label_name: Option<String>,
         operator_label_value: Option<String>,
         tunnel_client_default_image: String,
+        pia_canary_image: String,
     ) -> Self {
         Self {
             diagnostics: Arc::new(RwLock::new(Diagnostics::default())),
@@ -67,6 +71,7 @@ impl State {
             operator_label_name,
             operator_label_value,
             tunnel_client_default_image,
+            pia_canary_image,
         }
     }
 
diff --git a/src/controllers/restatecluster/controller.rs b/src/controllers/restatecluster/controller.rs
index 5c9c815..2a2a0ec 100644
--- a/src/controllers/restatecluster/controller.rs
+++ b/src/controllers/restatecluster/controller.rs
@@ -67,6 +67,8 @@ pub(super) struct Context {
     pub operator_label_name: Option<String>,
     // The value of the label named operator_label_name that will select the operator, needed to support the case where restate clusters need to be reached by the operator
     pub operator_label_value: Option<String>,
+    // The image to use for PIA canary pods
+    pub pia_canary_image: String,
     // Whether the EKS SecurityGroupPolicy CRD is installed
     pub security_group_policy_installed: bool,
     // Whether the SecretProviderClass CRD is installed
@@ -98,6 +100,7 @@ impl Context {
             operator_namespace: state.operator_namespace.clone(),
             operator_label_name: state.operator_label_name.clone(),
             operator_label_value: state.operator_label_value.clone(),
+            pia_canary_image: state.pia_canary_image.clone(),
             security_group_policy_installed,
             secret_provider_class_installed,
             diagnostics: state.diagnostics.clone(),
diff --git a/src/controllers/restatecluster/reconcilers/compute.rs b/src/controllers/restatecluster/reconcilers/compute.rs
index a1f3918..283d81e 100644
--- a/src/controllers/restatecluster/reconcilers/compute.rs
+++ b/src/controllers/restatecluster/reconcilers/compute.rs
@@ -551,6 +551,7 @@ pub async fn reconcile_compute(
                 name,
                 base_metadata,
                 spec.compute.tolerations.as_ref(),
+                &ctx.pia_canary_image,
                 &job_api,
                 &pod_api,
             )
@@ -745,6 +746,7 @@ async fn check_pia(
     namespace: &str,
     base_metadata: &ObjectMeta,
     tolerations: Option<&Vec<Toleration>>,
+    canary_image: &str,
     job_api: &Api<Job>,
     pod_api: &Api<Pod>,
 ) -> Result<(), Error> {
@@ -779,7 +781,7 @@ async fn check_pia(
                             service_account_name: Some("restate".into()),
                             containers: vec![Container {
                                 name: "canary".into(),
-                                image: Some("busybox:uclibc".into()),
+                                image: Some(canary_image.into()),
                                 command: Some(vec![
                                     "grep".into(),
                                     "-q".into(),
diff --git a/src/main.rs b/src/main.rs
index 481d866..af16ab6 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -45,6 +45,14 @@ struct Arguments {
         default_value = "ghcr.io/restatedev/restate-cloud-tunnel-client:0.5.0"
     )]
     tunnel_client_default_image: String,
+
+    #[arg(
+        long = "pia-canary-image",
+        env = "PIA_CANARY_IMAGE",
+        value_name = "IMAGE",
+        default_value = "busybox:uclibc"
+    )]
+    pia_canary_image: String,
 }
 
 #[get("/metrics")]
@@ -80,6 +88,7 @@ async fn main() -> anyhow::Result<()> {
         args.operator_label_name,
         args.operator_label_value,
         args.tunnel_client_default_image,
+        args.pia_canary_image,
     );
 
     let client = Client::try_default()