feat: Add opt-in upgrade_drain_timeout to drain control nodes before an upgrade

config/crd/bases/awx.ansible.com_awxs.yaml

@@ -209,6 +209,11 @@ spec:
                 description: Optional duration in seconds pods needs to terminate gracefully
                 type: integer
                 format: int32
+              upgrade_drain_timeout:
+                description: Seconds to wait for the control nodes' running jobs to finish before rolling to a new application version. 0 does not wait.
+                type: integer
+                format: int32
+                default: 0
               service_labels:
                 description: Additional labels to apply to the service
                 type: string

docs/user-guide/advanced-configuration/pods-termination-grace-period.md

@@ -29,11 +29,33 @@ following:
 
 One may want to set this value to the maximum duration they accept to wait for
 the affected Jobs to finish. Keeping in mind that such finishing jobs may
-increase Pods termination time in such situations as `kubectl rollout restart`,
-AWX upgrade by the operator, or Kubernetes [API-initiated
+increase Pods termination time in such situations as `kubectl rollout restart`
+or Kubernetes [API-initiated
 evictions](https://kubernetes.io/docs/concepts/scheduling-eviction/api-eviction/).
 
-
-| Name                             | Description                                                     | Default |
-| -------------------------------- | --------------------------------------------------------------- | ------- |
-| termination_grace_period_seconds | Optional duration in seconds pods needs to terminate gracefully | not set |
+#### Upgrades
+
+`termination_grace_period_seconds` does not cover an upgrade to a new AWX
+version on its own. Once the first Pod running the new version registers itself
+as an instance, every control node still running the old version stops its own
+services, because AWX shuts a node down when it sees a peer reporting a higher
+version. That happens from inside the container, so it preempts the `PreStop`
+hook, and the jobs the hook was waiting for fail with `Task was canceled due to
+receiving a shutdown signal.`
+
+Set `upgrade_drain_timeout` to also wait for those jobs *before* the new version
+is rolled out. When the application image changes, the operator disables each
+control node, so that new jobs stay pending rather than being scheduled onto a
+Pod that is about to be replaced, and waits for the jobs already running there
+to finish. Once they have, or once the timeout expires, the deployments are
+applied as usual. The wait happens before the schema migration runs, so the
+draining nodes are never running old code against a migrated database.
+
+A drain that reaches its timeout does not abort the upgrade: it proceeds, and
+any jobs still running are failed by the rollout as they would be without this
+setting.
+
+| Name                             | Description                                                                                                                                            | Default |
+| -------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------ | ------- |
+| termination_grace_period_seconds | Optional duration in seconds pods needs to terminate gracefully                                                                                          | not set |
+| upgrade_drain_timeout            | Seconds to wait for the control nodes' running jobs to finish before rolling the deployments to a new application version. 0 does not wait.               | 0       |

molecule/default/tasks/upgrade_drain_test.yml

@@ -0,0 +1,55 @@
+---
+# The molecule AWX spec sets upgrade_drain_timeout, but the application image
+# does not change between reconciles, so the drain must not run: an ordinary
+# reconcile has to leave every control node enabled and accepting jobs.
+- name: Validate upgrade_drain_timeout
+  tags:
+    - upgrade_drain
+  block:
+    - name: Get the AWX resource
+      k8s_info:
+        namespace: '{{ namespace }}'
+        api_version: awx.ansible.com/v1beta1
+        kind: AWX
+        name: example-awx
+      register: drain_awx
+
+    - name: Assert the spec carries upgrade_drain_timeout
+      ansible.builtin.assert:
+        that:
+          - drain_awx.resources | length == 1
+          - drain_awx.resources[0].spec.upgrade_drain_timeout | int > 0
+        fail_msg: 'upgrade_drain_timeout is not set on the AWX resource'
+
+    - name: Get the web pod
+      k8s_info:
+        namespace: '{{ namespace }}'
+        api_version: v1
+        kind: Pod
+        label_selectors:
+          - 'app.kubernetes.io/name=example-awx-web'
+        field_selectors:
+          - status.phase=Running
+      register: drain_web_pods
+
+    - name: Assert a web pod is running
+      ansible.builtin.assert:
+        that:
+          - drain_web_pods.resources | length
+        fail_msg: 'no running web pod to query the instances from'
+
+    - name: List the instances
+      k8s_exec:
+        namespace: '{{ namespace }}'
+        pod: '{{ drain_web_pods.resources[0].metadata.name }}'
+        container: example-awx-web
+        command: awx-manage list_instances
+      register: drain_instances
+      changed_when: false
+
+    # list_instances marks a disabled instance with a [DISABLED] prefix.
+    - name: Assert the reconcile left every control node enabled
+      ansible.builtin.assert:
+        that:
+          - "'[DISABLED]' not in drain_instances.stdout"
+        fail_msg: 'a reconcile without an image change disabled a control node: {{ drain_instances.stdout }}'

molecule/default/templates/awx_cr_molecule.yml.j2

@@ -38,6 +38,7 @@ spec:
       memory: 16M
   no_log: false
   restricted_security_context: true
+  upgrade_drain_timeout: 60
   postgres_resource_requirements: {}
   redis_resource_requirements: {}
   additional_labels:

molecule/default/verify.yml

@@ -18,6 +18,7 @@
             - tasks/awx_test.yml
             - tasks/awx_replicas_test.yml
             - tasks/restricted_security_context_test.yml
+            - tasks/upgrade_drain_test.yml
           tags:
             - always
       rescue:

roles/installer/defaults/main.yml

@@ -460,6 +460,11 @@ postgres_security_context_settings: {}
 # keeps existing behavior. task_privileged and redis_capabilities still apply.
 restricted_security_context: false
 
+# Seconds to wait for the control nodes' running jobs to finish before rolling
+# the deployments to a new application version. Opt-in; 0 keeps existing
+# behavior. See roles/installer/tasks/drain_control_nodes.yml.
+upgrade_drain_timeout: 0
+
 # Set no_log settings on certain tasks
 no_log: true
 

roles/installer/tasks/drain_control_nodes.yml

@@ -0,0 +1,101 @@
+---
+# Take the control nodes out of service and let their running jobs finish before
+# the deployments roll to a new application version.
+#
+# termination_grace_period_seconds cannot cover an upgrade on its own. As soon as
+# the first pod running the new version registers itself, every control node still
+# on the old version stops its own services (the version check in AWX's
+# cluster_node_heartbeat), which cuts short the preStop drain and fails the very
+# jobs that drain was waiting for. Draining has to happen while the cluster is
+# still wholly on the old version, which is here: the deployments have not been
+# applied yet and the schema migration runs later still.
+
+- name: Get the current task deployment
+  k8s_info:
+    api_version: apps/v1
+    kind: Deployment
+    namespace: '{{ ansible_operator_meta.namespace }}'
+    name: '{{ ansible_operator_meta.name }}-task'
+  register: _task_deployment
+
+- name: Set the currently deployed application image
+  set_fact:
+    _deployed_image: >-
+      {{ _task_deployment['resources'][0]['spec']['template']['spec']['containers']
+         | selectattr('name', 'equalto', ansible_operator_meta.name ~ '-task')
+         | map(attribute='image') | first | default('') }}
+  when: _task_deployment['resources'] | length > 0
+
+# Nothing to drain on a fresh install, when the application image is unchanged
+# (every other reconcile), or when no web pod is available to run awx-manage in.
+- name: Drain the control nodes
+  when:
+    - _deployed_image | default('') | length > 0
+    - _deployed_image != _image
+    - awx_web_pod_name | length > 0
+  block:
+    - name: Get the running control node pods
+      k8s_info:
+        api_version: v1
+        kind: Pod
+        namespace: '{{ ansible_operator_meta.namespace }}'
+        label_selectors:
+          - "app.kubernetes.io/name={{ ansible_operator_meta.name }}-task"
+          - "app.kubernetes.io/managed-by={{ deployment_type }}-operator"
+          - "app.kubernetes.io/component={{ deployment_type }}"
+        field_selectors:
+          - status.phase=Running
+      register: _control_pods
+
+    - name: Set the control node hostnames
+      set_fact:
+        _control_hostnames: >-
+          {{ _control_pods['resources']
+             | rejectattr('metadata.deletionTimestamp', 'defined')
+             | map(attribute='metadata.name') | list }}
+        _drain_retries: >-
+          {{ [1, ((upgrade_drain_timeout | int) / 10) | round(0, 'ceil') | int] | max }}
+
+    # Disable all of them before waiting on any of them: an instance that is still
+    # enabled keeps accepting work, so draining them one at a time would not
+    # converge. Disabling sets the instance capacity to zero, so jobs submitted
+    # from here on stay pending and run once the new pods register.
+    - name: Stop the control nodes accepting new jobs
+      k8s_exec:
+        namespace: '{{ ansible_operator_meta.namespace }}'
+        pod: '{{ awx_web_pod_name }}'
+        container: '{{ ansible_operator_meta.name }}-web'
+        command: awx-manage disable_instance --hostname {{ item }}
+      loop: '{{ _control_hostnames }}'
+      failed_when: false
+
+    # --retry=1 turns the command into a probe: it exits non-zero while the
+    # instance still has running or waiting jobs. Retrying from here keeps the
+    # k8s_exec calls short instead of holding one exec stream open for the whole
+    # drain.
+    - name: Wait for the control nodes' running jobs to finish
+      k8s_exec:
+        namespace: '{{ ansible_operator_meta.namespace }}'
+        pod: '{{ awx_web_pod_name }}'
+        container: '{{ ansible_operator_meta.name }}-web'
+        command: awx-manage disable_instance --hostname {{ item }} --wait --retry=1 --retry_sleep=1
+      loop: '{{ _control_hostnames }}'
+      register: _drain_result
+      until: _drain_result.return_code | default(1) == 0
+      retries: '{{ _drain_retries | int }}'
+      delay: 10
+      changed_when: false
+      ignore_errors: true
+
+    # A drain that never finishes must not block the upgrade, so this is a
+    # warning rather than a failure: the jobs still running on those nodes will
+    # fail as they would without this feature.
+    - name: Report control nodes that were still running jobs
+      debug:
+        msg: >-
+          {{ item.item }} still had running jobs after upgrade_drain_timeout
+          ({{ upgrade_drain_timeout }}s); continuing with the upgrade.
+      loop: '{{ _drain_result.results | default([]) }}'
+      loop_control:
+        label: '{{ item.item }}'
+      when: item.return_code | default(1) != 0

roles/installer/tasks/resources_configuration.yml

@@ -245,6 +245,10 @@
   set_fact:
     _redis_image: "{{ _custom_redis_image | default(lookup('env', 'RELATED_IMAGE_AWX_REDIS')) | default(_default_redis_image, true) }}"
 
+- name: Let the control nodes' running jobs finish before rolling a new version
+  include_tasks: drain_control_nodes.yml
+  when: upgrade_drain_timeout | int > 0
+
 - name: Apply deployment resources
   k8s:
     apply: yes