Support restart task for StatefulSet with updateStrategy: OnDelete

CHANGELOG.md

@@ -1,5 +1,9 @@
 ## next
 
+*Enhancements*
+
+- Support restart task for stateful sets that use the `OnDelete` strategy [#840](https://github.com/Shopify/krane/pull/840)
+
 ## 2.4.0
 
 *Enhancements*

lib/krane/kubernetes_resource/stateful_set.rb

@@ -27,13 +27,11 @@ def deploy_succeeded?
                        "Consider switching to rollingUpdate.")
           @success_assumption_warning_shown = true
         end
-        true
-      else
-        observed_generation == current_generation &&
-        status_data['currentRevision'] == status_data['updateRevision'] &&
-        desired_replicas == status_data['readyReplicas'].to_i &&
-        desired_replicas == status_data['currentReplicas'].to_i
       end
+      observed_generation == current_generation &&
+      status_data['currentRevision'] == status_data['updateRevision'] &&
+      desired_replicas == status_data['readyReplicas'].to_i &&
+      desired_replicas == status_data['currentReplicas'].to_i
     end
 
     def deploy_failed?

lib/krane/restart_task.rb

@@ -69,10 +69,10 @@ def run!(deployments: [], statefulsets: [], daemonsets: [], selector: nil, verif
       deployments, statefulsets, daemonsets = identify_target_workloads(deployments, statefulsets,
                                                 daemonsets, selector: selector)
 
-      @logger.phase_heading("Triggering restart by annotating pod template #{RESTART_TRIGGER_ANNOTATION} annotation")
-      patch_kubeclient_deployments(deployments)
-      patch_kubeclient_statefulsets(statefulsets)
-      patch_kubeclient_daemonsets(daemonsets)
+      @logger.phase_heading("Triggering restart")
+      restart_deployments!(deployments)
+      restart_statefulsets!(statefulsets)
+      restart_daemonsets!(daemonsets)
 
       if verify_result
         @logger.phase_heading("Waiting for rollout")
@@ -210,7 +210,15 @@ def patch_daemonset_with_restart(record)
       apps_v1_kubeclient.patch_daemon_set(record.metadata.name, build_patch_payload(record), @namespace)
     end
 
-    def patch_kubeclient_deployments(deployments)
+    def delete_statefulset_pods(record)
+      pods = kubeclient.get_pods(namespace: record.metadata.namespace)
+      pods.select! do |pod|
+        pod.metadata.ownerReferences.find { |ref| ref.uid == record.metadata.uid }
+      end
+      pods.each { |pod| kubeclient.delete_pod(pod.metadata.name, pod.metadata.namespace) }
+    end
+
+    def restart_deployments!(deployments)
       deployments.each do |record|
         begin
           patch_deployment_with_restart(record)
@@ -221,18 +229,23 @@ def patch_kubeclient_deployments(deployments)
       end
     end
 
-    def patch_kubeclient_statefulsets(statefulsets)
+    def restart_statefulsets!(statefulsets)
       statefulsets.each do |record|
-        begin
+        @logger.info("Triggered `StatefulSet/#{record.metadata.name}` restart")
+        if record.spec.updateStrategy&.type == "OnDelete"
+          @logger.info("`StatefulSet/#{record.metadata.name}` has updateStrategy: OnDelete," \
+            " Restarting by forcefully deleting child pods"
+          )
+          delete_statefulset_pods(record)
+        else
           patch_statefulset_with_restart(record)
-          @logger.info("Triggered `StatefulSet/#{record.metadata.name}` restart")
-        rescue Kubeclient::HttpError => e
-          raise RestartAPIError.new(record.metadata.name, e.message)
         end
+      rescue Kubeclient::HttpError => e
+        raise RestartAPIError.new(record.metadata.name, e.message)
       end
     end
 
-    def patch_kubeclient_daemonsets(daemonsets)
+    def restart_daemonsets!(daemonsets)
       daemonsets.each do |record|
         begin
           patch_daemonset_with_restart(record)

test/integration/restart_task_test.rb

@@ -35,6 +35,35 @@ def test_restart_by_annotation
     refute(fetch_restarted_at("redis"), "no restart annotation env on fresh deployment")
   end
 
+  def test_restart_statefulset_on_delete_restarts_child_pods
+    result = deploy_fixtures("hello-cloud", subset: "stateful_set.yml") do |fixtures|
+      statefulset = fixtures["stateful_set.yml"]["StatefulSet"].first
+      statefulset["spec"]["updateStrategy"] = { "type" => "OnDelete" }
+    end
+    assert_deploy_success(result)
+    before_pods = kubeclient.get_pods(namespace: @namespace, label_selector: "name=stateful-busybox").map do |p|
+      p.metadata.uid
+    end
+
+    restart = build_restart_task
+    assert_restart_success(restart.perform)
+    after_pods = kubeclient.get_pods(namespace: @namespace, label_selector: "name=stateful-busybox").map do |p|
+      p.metadata.uid
+    end
+    refute_equal(before_pods.sort, after_pods.sort)
+
+    assert_logs_match_all([
+      "Configured to restart all workloads with the `shipit.shopify.io/restart` annotation",
+      "Triggered `StatefulSet/stateful-busybox` restart",
+      "`StatefulSet/stateful-busybox` has updateStrategy: OnDelete, Restarting by forcefully deleting child pods",
+      "Waiting for rollout",
+      "Result: SUCCESS",
+      "Successfully restarted 1 resource",
+      %r{StatefulSet/stateful-busybox.* 2 replicas},
+    ],
+      in_order: true)
+  end
+
   def test_restart_by_selector
     assert_deploy_success(deploy_fixtures("branched",
       bindings: { "branch" => "master" },