commit a944005a50cb94d88520af530397e35eec6c3108
parent d12459d1c6e7c5c0a23f857eb63b9cc98a0062e2
Author: Asher Morgan <asher@ashermorgan.net>
Date: Sun, 30 Aug 2026 19:26:25 -0700
Implement keep-alive messages
Diffstat:
2 files changed, 74 insertions(+), 49 deletions(-)
diff --git a/aoh/client.py b/aoh/client.py
@@ -20,12 +20,15 @@ PASSWORD_PROMPTS = {
'vault_password': 'Vault password: ',
}
+# Set 10s EXEC keep alive interval (see also TIMEOUT in aoh.py)
+KEEP_ALIVE_INTERVAL = 10
+
class ClientError(Exception):
"""Raised for miscellaneous client errors."""
-def exec(args):
+def exec(args, keep_alive_handler=None):
"""Run a command on the local host."""
try:
@@ -36,6 +39,14 @@ def exec(args):
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
)
+
+ while p.returncode is None:
+ try:
+ p.wait(KEEP_ALIVE_INTERVAL)
+ except subprocess.TimeoutExpired:
+ if keep_alive_handler:
+ keep_alive_handler()
+
stdout, stderr = p.communicate()
return {
'returncode': p.returncode,
@@ -103,33 +114,48 @@ def send_json_request(url, req_data=None, method='GET', cookies=''):
return status_code, headers, json.loads(data)
-def runner_loop(runner_url, cookies):
- """Main execution loop."""
+def runner_update(url, cookies, data, keep_alive=False):
+ """Submit runner updates and process core response fields."""
- req = {}
+ status, _, res = send_json_request(url, data, 'PUT', cookies)
- while True:
- status, _, res = send_json_request(runner_url, req, 'PUT', cookies)
+ if 'logs' in res:
+ # Strip duplicate password prompts
+ logs = res['logs']
+ while any(logs.startswith(prompt[:-2]) for prompt in
+ PASSWORD_PROMPTS.values()):
+ logs = logs.split('\n', 1)[1]
+
+ print(logs, end='')
+
+ if 'err' in res:
+ raise ClientError(res['err'])
+ elif status != 200:
+ raise ClientError(f'Server responded with {status}')
+ elif keep_alive and not res.get('keep-alive'):
+ raise ClientError('Received invalid keep-alive response')
+ elif res.get('finished'):
+ sys.exit(0)
- if 'logs' in res:
- # Strip duplicate password prompts
- logs = res['logs']
- while any(logs.startswith(prompt[:-2]) for prompt in
- PASSWORD_PROMPTS.values()):
- logs = logs.split('\n', 1)[1]
+ return res
- print(logs, end='')
- if 'err' in res:
- raise ClientError(res['err'])
- elif status != 200:
- raise ClientError(f'Server responded with {status}')
+def runner_loop(runner_url, cookies):
+ """Main execution loop."""
+
+ req = {}
- if res.get('finished'):
- return
+ while True:
+ res = runner_update(runner_url, cookies, req)
if 'exec' in res:
- req = exec(res['exec'])
+ keep_alive_handler = lambda res=res: runner_update(
+ runner_url,
+ cookies,
+ { 'id': res['id'], 'keep-alive': True },
+ keep_alive=True
+ )
+ req = exec(res['exec'], keep_alive_handler)
req['id'] = res['id']
elif 'put' in res:
req = put(res['put'])
diff --git a/aoh/connection_plugins/aoh.py b/aoh/connection_plugins/aoh.py
@@ -12,19 +12,6 @@ DOCUMENTATION = """
required: true
vars:
- name: ansible_aoh_dir
- aoh_timeout:
- description: The timeout for AoH responses, in seconds
- type: integer
- default: 300
- env:
- - name: ANSIBLE_TIMEOUT
- ini:
- - key: timeout
- section: defaults
- vars:
- - name: ansible_aoh_timeout
- cli:
- - name: timeout
"""
import base64
@@ -45,6 +32,9 @@ from ansible.utils.display import Display
display = Display()
+# Set 1m timeout (see also KEEP_ALIVE_INTERVAL in client.py)
+TIMEOUT = 60
+
class Connection(ConnectionBase):
"""Ansible-over-HTTP (AoH) connection."""
@@ -107,22 +97,31 @@ class Connection(ConnectionBase):
self.sendbuf.write(json.dumps(msg) + '\n')
self.sendbuf.flush()
- if not self.recvpoll.poll(self.get_option('aoh_timeout') * 1000):
- raise AnsibleConnectionFailure('Timed out waiting for AoH response')
-
- try:
- res = json.loads(self.recvbuf.readline())
- except json.decoder.JSONDecodeError:
- raise AnsibleError('Received corrupt AoH message')
-
- if not 'id' in msg:
- raise AnsibleError('Received AoH message without ID')
- elif res['id'] != msg['id']:
- raise AnsibleError('Received AoH message with unexpected ID')
- elif 'err' in res:
- raise AnsibleError(str(res['err']))
-
- return res
+ while True:
+ if not self.recvpoll.poll(TIMEOUT * 1000):
+ raise AnsibleConnectionFailure('Timed out waiting for AoH '
+ 'response')
+
+ try:
+ res = json.loads(self.recvbuf.readline())
+ except json.decoder.JSONDecodeError:
+ raise AnsibleError('Received corrupt AoH message')
+
+ if not 'id' in msg:
+ raise AnsibleError('Received AoH message without ID')
+ elif res['id'] != msg['id']:
+ raise AnsibleError('Received AoH message with unexpected ID')
+ elif 'err' in res:
+ raise AnsibleError(str(res['err']))
+
+ if not res.get('keep-alive'):
+ return res
+
+ self.sendbuf.write(json.dumps({
+ 'id': msg['id'],
+ 'keep-alive': True,
+ }) + '\n')
+ self.sendbuf.flush()
def exec_command(self, cmd: str, in_data: bytes | None = None,
