Add support for specifying restate pod tolerations (#11)

* Add support for specifying restate pod tolerations

* Generate pkl

---------

Co-authored-by: Jack Kleeman <jackkleeman@gmail.com>

charts/restate-operator-helm/templates/deployment.yaml

@@ -55,3 +55,7 @@ spec:
               port: http
             initialDelaySeconds: 5
             periodSeconds: 5
+      {{- if .Values.tolerations }}
+      tolerations:
+        {{- toYaml .Values.tolerations | nindent 8 }}
+      {{- end }}

crd/RestateCluster.pkl

@@ -73,6 +73,36 @@ class Compute {
   /// Compute Resources for the Restate container. More info:
   /// https://kubernetes.io/docs/concepts/configuration/manage-resources-containers/
   resources: ResourceRequirements?
+
+  /// If specified, the pod's tolerations.
+  tolerations: Listing<Toleration>?
+}
+
+/// The pod this Toleration is attached to tolerates any taint that matches the triple <key,value,effect>
+/// using the matching operator <operator>.
+class Toleration {
+  /// Effect indicates the taint effect to match. Empty means match all taint effects. When specified,
+  /// allowed values are NoSchedule, PreferNoSchedule and NoExecute.
+  effect: String?
+
+  /// Key is the taint key that the toleration applies to. Empty means match all taint keys. If the key
+  /// is empty, operator must be Exists; this combination means to match all values and all keys.
+  key: String?
+
+  /// Operator represents a key's relationship to the value. Valid operators are Exists and Equal.
+  /// Defaults to Equal. Exists is equivalent to wildcard for value, so that a pod can tolerate all
+  /// taints of a particular category.
+  operator: String?
+
+  /// TolerationSeconds represents the period of time the toleration (which must be of effect NoExecute,
+  /// otherwise this field is ignored) tolerates the taint. By default, it is not set, which means
+  /// tolerate the taint forever (do not evict). Zero and negative values will be treated as 0 (evict
+  /// immediately) by the system.
+  tolerationSeconds: Int?
+
+  /// Value is the taint value the toleration matches to. If the operator is Exists, the value should be
+  /// empty, otherwise just a regular string.
+  value: String?
 }
 
 /// Security configuration

crd/crd.yaml

@@ -205,6 +205,30 @@ spec:
                         description: 'Requests describes the minimum amount of compute resources required. If Requests is omitted for a container, it defaults to Limits if that is explicitly specified, otherwise to an implementation-defined value. Requests cannot exceed Limits. More info: https://kubernetes.io/docs/concepts/configuration/manage-resources-containers/'
                         type: object
                     type: object
+                  tolerations:
+                    description: If specified, the pod's tolerations.
+                    items:
+                      description: The pod this Toleration is attached to tolerates any taint that matches the triple <key,value,effect> using the matching operator <operator>.
+                      properties:
+                        effect:
+                          description: Effect indicates the taint effect to match. Empty means match all taint effects. When specified, allowed values are NoSchedule, PreferNoSchedule and NoExecute.
+                          type: string
+                        key:
+                          description: Key is the taint key that the toleration applies to. Empty means match all taint keys. If the key is empty, operator must be Exists; this combination means to match all values and all keys.
+                          type: string
+                        operator:
+                          description: Operator represents a key's relationship to the value. Valid operators are Exists and Equal. Defaults to Equal. Exists is equivalent to wildcard for value, so that a pod can tolerate all taints of a particular category.
+                          type: string
+                        tolerationSeconds:
+                          description: TolerationSeconds represents the period of time the toleration (which must be of effect NoExecute, otherwise this field is ignored) tolerates the taint. By default, it is not set, which means tolerate the taint forever (do not evict). Zero and negative values will be treated as 0 (evict immediately) by the system.
+                          format: int64
+                          type: integer
+                        value:
+                          description: Value is the taint value the toleration matches to. If the operator is Exists, the value should be empty, otherwise just a regular string.
+                          type: string
+                      type: object
+                    nullable: true
+                    type: array
                 required:
                 - image
                 type: object

src/controller.rs

@@ -11,7 +11,7 @@ use k8s_openapi::api::apps::v1::StatefulSet;
 use k8s_openapi::api::batch::v1::Job;
 use k8s_openapi::api::core::v1::{
     ConfigMap, EnvVar, Namespace, PersistentVolumeClaim, PodDNSConfig, ResourceRequirements,
-    Service, ServiceAccount,
+    Service, ServiceAccount, Toleration,
 };
 use k8s_openapi::api::networking::v1;
 use k8s_openapi::api::networking::v1::{NetworkPolicy, NetworkPolicyPeer, NetworkPolicyPort};
@@ -191,6 +191,8 @@ pub struct RestateClusterCompute {
     pub dns_config: Option<PodDNSConfig>,
     /// Set DNS policy for the pod. Defaults to "ClusterFirst". Valid values are 'ClusterFirstWithHostNet', 'ClusterFirst', 'Default' or 'None'. DNS parameters given in DNSConfig will be merged with the policy selected with DNSPolicy.
     pub dns_policy: Option<String>,
+    /// If specified, the pod's tolerations.
+    pub tolerations: Option<Vec<Toleration>>,
 }
 
 fn env_schema(g: &mut schemars::gen::SchemaGenerator) -> Schema {
@@ -499,6 +501,7 @@ impl RestateCluster {
 
         Ok(())
     }
+
     async fn reconcile_status(&self, ctx: Arc<Context>) -> Result<Action> {
         let rcs: Api<RestateCluster> = Api::all(ctx.client.clone());
 

src/main.rs

@@ -44,7 +44,7 @@ async fn main() -> anyhow::Result<()> {
 
     let args: Arguments = Arguments::parse();
 
-    // Initiatilize Kubernetes controller state
+    // Initialize Kubernetes controller state
     let state = State::default().with_aws_pod_identity_association_cluster(
         args.aws_pod_identity_association_cluster
             .and_then(|s| s.to_str().map(|s| s.to_string())),

src/reconcilers/compute.rs

@@ -332,6 +332,7 @@ fn restate_statefulset(
                     service_account_name: Some("restate".into()),
                     termination_grace_period_seconds: Some(60),
                     volumes: Some(volumes),
+                    tolerations: spec.compute.tolerations.clone(),
                     ..Default::default()
                 }),
             },