Handling the job cancel in case of timeouts

providers/amazon/src/airflow/providers/amazon/aws/triggers/base.py

@@ -84,6 +84,7 @@ def __init__(
         region_name: str | None = None,
         verify: bool | str | None = None,
         botocore_config: dict | None = None,
+        cancel_waiter_names: list[str] | None = None,
     ):
         super().__init__()
         # parameters that should be hardcoded in the child's implem
@@ -105,6 +106,7 @@ def __init__(
         self.region_name = region_name
         self.verify = verify
         self.botocore_config = botocore_config
+        self.cancel_waiter_names = cancel_waiter_names
 
     def serialize(self) -> tuple[str, dict[str, Any]]:
         # here we put together the "common" params,
@@ -159,3 +161,25 @@ async def run(self) -> AsyncIterator[TriggerEvent]:
                 self.status_queries,
             )
             yield TriggerEvent({"status": "success", self.return_key: self.return_value})
+
+    async def cleanup(self):
+        hook = self.hook()
+        async with await hook.get_async_conn() as client:
+            if not self.cancel_waiter_names:
+                return
+            for waiter_name in self.cancel_waiter_names:
+                waiter = hook.get_waiter(
+                    waiter_name,
+                    deferrable=True,
+                    client=client,
+                    config_overrides=self.waiter_config_overrides,
+                )
+                await async_wait(
+                    waiter,
+                    self.waiter_delay,
+                    self.attempts,
+                    self.waiter_args,
+                    self.failure_message,
+                    self.status_message,
+                    self.status_queries,
+                )

providers/amazon/src/airflow/providers/amazon/aws/triggers/emr.py

@@ -174,6 +174,7 @@ def __init__(
             waiter_delay=waiter_delay,
             waiter_max_attempts=waiter_max_attempts,
             aws_conn_id=aws_conn_id,
+            cancel_waiter_names=["container_job_cancel", "container_job_complete"],
         )
 
     def hook(self) -> AwsGenericHook:

providers/amazon/src/airflow/providers/amazon/aws/waiters/emr-containers.json

@@ -25,6 +25,18 @@
                     "state": "failure"
                 }
             ]
+        },
+        "container_job_cancel": {
+            "operation": "CancelJobRun",
+            "delay": 30,
+            "maxAttempts": 60,
+            "acceptors": [
+                {
+                    "matcher": "status",
+                    "expected": "200",
+                    "state": "success"
+                }
+            ]
         }
     }
 }