Speed up the default preStop hook to allow for graceful shutdown

paasta_tools/kubernetes_tools.py

@@ -1434,7 +1434,9 @@ def get_kubernetes_containers(
             env=self.get_container_env(),
             resources=self.get_resource_requirements(),
             lifecycle=V1Lifecycle(
-                pre_stop=self.get_kubernetes_container_termination_action()
+                pre_stop=self.get_kubernetes_container_termination_action(
+                    service_namespace_config
+                )
             ),
             name=self.get_sanitised_instance_name(),
             liveness_probe=self.get_liveness_probe(service_namespace_config),
@@ -1464,17 +1466,36 @@ def get_readiness_probe(
         else:
             return self.get_liveness_probe(service_namespace_config)
 
-    def get_kubernetes_container_termination_action(self) -> V1Handler:
+    def get_kubernetes_container_termination_action(
+        self, service_namespace_config: ServiceNamespaceConfig
+    ) -> V1Handler:
         command = self.config_dict.get("lifecycle", KubeLifecycleDict({})).get(
             "pre_stop_command", []
         )
         # default pre stop hook for the container
         if not command:
-            return V1Handler(
-                _exec=V1ExecAction(
-                    command=["/bin/sh", "-c", f"sleep {DEFAULT_PRESTOP_SLEEP_SECONDS}"]
+            # if the service is in smartstack, we want to sleep for a while to allow requests to drain
+            if service_namespace_config.is_in_smartstack():
+                return V1Handler(
+                    _exec=V1ExecAction(
+                        command=[
+                            "/bin/sh",
+                            "-c",
+                            f"sleep {DEFAULT_PRESTOP_SLEEP_SECONDS}",
+                        ]
+                    )
+                )
+            # otherwise, we want to immediately send a SIGTERM to the container
+            # so that it can gracefully shutdown if configured to do so
+            else:
+                return V1Handler(
+                    _exec=V1ExecAction(
+                        # NOTE: we could also probably look into seeing what happens if we return None
+                        # instead of a V1Handler, but in the meantime, having a funny message as the pre-stop
+                        # hook doesn't seem particularly harmful
+                        command=["/bin/sh", "-c", f"echo Goodbye, cruel world."]
+                    )
                 )
-            )
         if isinstance(command, str):
             command = [command]
         return V1Handler(_exec=V1ExecAction(command=command))

tests/test_kubernetes_tools.py

@@ -955,6 +955,7 @@ def test_get_kubernetes_containers(self, prometheus_port, expected_ports):
             service_namespace_config = mock.Mock()
             service_namespace_config.get_healthcheck_mode.return_value = "http"
             service_namespace_config.get_healthcheck_uri.return_value = "/status"
+            service_namespace_config.is_in_smartstack.return_value = True
             assert (
                 self.deployment.get_kubernetes_containers(
                     docker_volumes=mock_docker_volumes,
@@ -2185,24 +2186,59 @@ def test_get_pod_anti_affinity(self, anti_affinity, expected):
         assert self.deployment.get_pod_anti_affinity() == expected_affinity
 
     @pytest.mark.parametrize(
-        "termination_action,expected",
+        "termination_action,expected,service_namespace_config",
         [
-            (None, ["/bin/sh", "-c", "sleep 30"]),  # no termination action
-            ("", ["/bin/sh", "-c", "sleep 30"]),  # empty termination action
-            ([], ["/bin/sh", "-c", "sleep 30"]),  # empty termination action
-            ("/bin/no-args", ["/bin/no-args"]),  # no args command
-            (["/bin/bash", "cmd.sh"], ["/bin/bash", "cmd.sh"]),  # no args command
+            (
+                None,
+                ["/bin/sh", "-c", "echo Goodbye, cruel world."],
+                ServiceNamespaceConfig({}),
+            ),  # no termination action
+            (
+                "",
+                ["/bin/sh", "-c", "echo Goodbye, cruel world."],
+                ServiceNamespaceConfig({}),
+            ),  # empty termination action
+            (
+                [],
+                ["/bin/sh", "-c", "echo Goodbye, cruel world."],
+                ServiceNamespaceConfig({}),
+            ),  # empty termination action
+            (
+                None,
+                ["/bin/sh", "-c", "sleep 30"],
+                ServiceNamespaceConfig({"proxy_port": 6666}),
+            ),  # no termination action, mesh-registered
+            (
+                "",
+                ["/bin/sh", "-c", "sleep 30"],
+                ServiceNamespaceConfig({"proxy_port": 6666}),
+            ),  # empty termination action, mesh-registered
+            (
+                [],
+                ["/bin/sh", "-c", "sleep 30"],
+                ServiceNamespaceConfig({"proxy_port": 6666}),
+            ),  # empty termination action, mesh-registered
+            ("/bin/no-args", ["/bin/no-args"], {}),  # no args command
+            (["/bin/bash", "cmd.sh"], ["/bin/bash", "cmd.sh"], {}),  # no args command
         ],
     )
     def test_kubernetes_container_termination_action(
-        self, termination_action, expected
+        self,
+        termination_action,
+        expected,
+        service_namespace_config,
     ):
         if termination_action:
             self.deployment.config_dict["lifecycle"] = {
                 "pre_stop_command": termination_action
             }
         handler = V1Handler(_exec=V1ExecAction(command=expected))
-        assert self.deployment.get_kubernetes_container_termination_action() == handler
+        assert (
+            self.deployment.get_kubernetes_container_termination_action(
+                service_namespace_config
+            )
+            == handler
+        )
 
     @pytest.mark.parametrize(
         "whitelist,blacklist,expected",
@@ -4274,7 +4310,7 @@ def test_warning_big_bounce():
             job_config.format_kubernetes_app().spec.template.metadata.labels[
                 "paasta.yelp.com/config_sha"
             ]
-            == "config5abf6b07"
+            == "config7d73e761"
         ), "If this fails, just change the constant in this test, but be aware that deploying this change will cause every service to bounce!"